Add tolerance for train_loss check in test_fsdp.py

imangohari1 · Oct 4, 2024 · e625fce · e625fce
1 parent 8ed336a
commit e625fce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_fsdp_examples.py b/tests/test_fsdp_examples.py
@@ -154,7 +154,7 @@ def _test_fsdp(
         if model_name == "bert-base-uncased":
             assert results["eval_f1"] >= ACCURACY_PERF_FACTOR * baseline_acc
         else:
-            assert results["train_loss"] <= baseline_acc
+            assert results["train_loss"] <= (2 - ACCURACY_PERF_FACTOR) * baseline_acc
 
 
 @pytest.mark.parametrize(