fix

huggingface · Feb 10, 2024 · 63ca0d2 · 63ca0d2
1 parent e3c3d11
commit 63ca0d2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/tests/helpers/utils.py b/tests/helpers/utils.py
@@ -216,7 +216,7 @@ def test_something():
     return func_wrapper
 
 
-def rerun_on_exception(exception_type: Exception = Exception, pattern: str = None, max_try: int = 5) -> Callable:
+def rerun_on_exception(exception_type: Exception = Exception, pattern: str = None, max_try: int = 10) -> Callable:
     """
     A decorator on a function to re-run when an exception occurs.
 

diff --git a/tests/test_clip_grads.py b/tests/test_clip_grads.py
@@ -422,8 +422,8 @@ def _test_clip_grads_tied_weights(parallel_context: ParallelContext, norm_type:
         norm_type=norm_type,
     )
     ref_total_norm = torch.nn.utils.clip_grad_norm_([ref_weight, ref_bias], max_norm=1.0, norm_type=norm_type)
-    assert total_norm.dim() == 0, f"total_norm should be a scalar. Got {total_norm}"
 
+    assert total_norm.dim() == 0, f"total_norm should be a scalar. Got {total_norm}"
     # Check that the gradients have changed
     assert not torch.allclose(old_grad, weight.grad), "Gradients should have changed after clipping"