Patchfix infinite loop (#335)

huggingface · Apr 29, 2022 · b3691db · b3691db
1 parent c665779
commit b3691db
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/accelerate/memory_utils.py b/src/accelerate/memory_utils.py
@@ -82,5 +82,7 @@ def decorator(*args, **kwargs):
                     gc.collect()
                     torch.cuda.empty_cache()
                     batch_size //= 2
+                else:
+                    raise
 
     return decorator
diff --git a/tests/test_memory_utils.py b/tests/test_memory_utils.py
@@ -80,3 +80,12 @@ def mock_training_loop_function(batch_size, arg1, arg2):
             mock_training_loop_function(128, "hello", "world")
             self.assertIn("Batch size was passed into `f`", cm.exception.args[0])
             self.assertIn("`f(arg1='hello', arg2='world')", cm.exception.args[0])
+
+    def test_any_other_error(self):
+        @find_executable_batch_size(starting_batch_size=16)
+        def mock_training_loop_function(batch_size):
+            raise ValueError("Oops, we had an error!")
+
+        with self.assertRaises(ValueError) as cm:
+            mock_training_loop_function()
+            self.assertIn("Oops, we had an error!", cm.exception.args[0])