Fix check (#6798)

Signed-off-by: MaximumEntropy <[email protected]>
NVIDIA · Jun 2, 2023 · 9db6009 · 9db6009
1 parent 5428a97
commit 9db6009
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/data/language_modeling/megatron/gpt_dataset.py b/nemo/collections/nlp/data/language_modeling/megatron/gpt_dataset.py
@@ -601,7 +601,7 @@ def _build_index_mappings(
                 last_epoch_num_samples = num_samples - num_samples_from_epochs_minus_one
                 assert last_epoch_num_samples >= 0, 'last epoch number of samples should be non-negative.'
                 num_samples_per_epoch = (tokens_per_epoch - add_extra_token) // seq_length
-                assert last_epoch_num_samples < (
+                assert last_epoch_num_samples <= (
                     num_samples_per_epoch + 1
                 ), 'last epoch number of samples exceeded max value.'
                 # If we have less than 80% of the samples for the last epoch,