Fixed bug with extra batch being run, unwrapped combined operators for clarity (#890)

Author: kylesayrs

src/sparseml/pytorch/utils/module.py

@@ -693,6 +693,9 @@ def run(
         epoch_timer = time.time()
 
         for batch, data in data_iter:
+            if 0 < max_steps and batch >= max_steps:
+                break
+
             step_timer = time.time()
             batch_size = self._run_funcs.batch_size(data)  # type: int
 
@@ -751,9 +754,6 @@ def run(
             if results is not None:
                 results.append(batch_results, batch_size)
 
-            if 0 < max_steps <= batch:
-                break
-
         should_log = self._loggers and self._log_summary and results
         log_step = counter  # log under the counter step for the summaries