Fix training loop

huggingface · Feb 6, 2025 · f0bc82e · f0bc82e
1 parent 34d59f4
commit f0bc82e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/optimum/neuron/trainers.py b/optimum/neuron/trainers.py
@@ -1114,6 +1114,11 @@ def _inner_training_loop(
                         if is_torch_xla_available():
                             xm.mark_step()
                         break
+                # We also need to break out of the nested loop
+                if self.control.should_epoch_stop or self.control.should_training_stop:
+                    if is_torch_xla_available():
+                        xm.mark_step()
+                    break
             if step < 0:
                 if is_main_worker():
                     logger.warning(