Add remaining stuffs

pomonam · Jun 22, 2024 · cce322b · cce322b
1 parent 523e7d8
commit cce322b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/modules/test_per_sample_gradients.py b/tests/modules/test_per_sample_gradients.py
@@ -179,7 +179,6 @@ def test_for_loop_per_sample_gradient_equivalence(
     "test_name",
     [
         "mlp",
-        "repeated_mlp",
         "conv",
         "conv_bn",
         "bert",
@@ -242,6 +241,7 @@ def test_mean_gradient_equivalence(
                 module_gradients[module.name] = module.get_factor(factor_name=PRECONDITIONED_GRADIENT_NAME)
 
         per_sample_gradients.append(module_gradients)
+        set_mode(model, ModuleMode.PRECONDITION_GRADIENT, keep_factors=False)
 
     summed_gradients = []
     for i in range(num_batches):
@@ -291,7 +291,6 @@ def test_mean_gradient_equivalence(
     "test_name",
     [
         "mlp",
-        "repeated_mlp",
         "conv",
         "gpt",
     ],

diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py
@@ -112,6 +112,7 @@ def test_default_score_arguments() -> None:
     assert score_args.cached_activation_cpu_offload is False
     assert score_args.distributed_sync_steps == 1000
     assert score_args.amp_dtype is None
+    assert score_args.shared_parameters_exist is False
 
     assert score_args.data_partition_size == 1
     assert score_args.module_partition_size == 1