Update source/tests/pt/test_multitask.py

Co-authored-by: Duo <[email protected]> Signed-off-by: Han Wang <[email protected]>
deepmodeling · Apr 6, 2024 · 4714684 · 4714684
1 parent 0dae3c9
commit 4714684
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py
@@ -138,7 +138,7 @@ def test_multitask_train(self):
                     multi_state_dict[state_key.replace("model_3", "model_2")],
                     multi_state_dict_finetuned[state_key],
                 )
-            elif "model_4" in state_key and "atomic_model" not in state_key:
+            elif "model_4" in state_key and "fitting_net" not in state_key and "out_bias" not in state_key:
                 torch.testing.assert_close(
                     multi_state_dict[state_key.replace("model_4", "model_2")],
                     multi_state_dict_finetuned[state_key],