fix test

hiyouga · Oct 22, 2024 · 74a79cc · 74a79cc
1 parent d155b70
commit 74a79cc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/llamafactory/train/test_utils.py b/src/llamafactory/train/test_utils.py
@@ -37,9 +37,9 @@ def compare_model(model_a: "torch.nn.Module", model_b: "torch.nn.Module", diff_k
     assert set(state_dict_a.keys()) == set(state_dict_b.keys())
     for name in state_dict_a.keys():
         if any(key in name for key in diff_keys):
-            assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-3, atol=1e-4) is False
+            assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-2, atol=1e-3) is False
         else:
-            assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-3, atol=1e-4) is True
+            assert torch.allclose(state_dict_a[name], state_dict_b[name], rtol=1e-2, atol=1e-3) is True
 
 
 def check_lora_model(model: "LoraModel") -> Tuple[Set[str], Set[str]]: