Fixed minor float check bug in compute_num_steps_per_epoch (#235)

Signed-off-by: Daniel Egert <degert@nvidia.com>
NVIDIA · Jul 12, 2024 · 23752e7 · 23752e7
1 parent 2a454f9
commit 23752e7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo_aligner/utils/trainer_utils.py b/nemo_aligner/utils/trainer_utils.py
@@ -29,7 +29,7 @@ def compute_num_steps_per_epoch(
 
     num_steps_per_epoch = sampler.total_samples // sampler.global_batch_size
 
-    if limit_train_batches is None or limit_train_batches > 1.0:
+    if limit_train_batches is None or (isinstance(limit_train_batches, float) and limit_train_batches > 1.0):
         limit_train_batches = 1.0
 
     if limit_train_batches >= 0: