Fix samples mapping bug

e3c57c82 · Neel Kant · bf599e86 · e3c57c82
Commit e3c57c82 authored Jun 08, 2020 by Neel Kant
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

megatron/data/realm_dataset_utils.py megatron/data/realm_dataset_utils.py +2 -2

No files found.
--- a/megatron/data/realm_dataset_utils.py
+++ b/megatron/data/realm_dataset_utils.py
@@ -221,8 +221,8 @@ def get_block_samples_mapping(block_dataset, title_dataset, data_prefix, num_epo
    # parallel case
    counts = torch.cuda.LongTensor([1])
    torch.distributed.all_reduce(counts, group=mpu.get_data_parallel_group())
-    #assert counts[0].item() == torch.distributed.get_world_size(
-    #    group=mpu.get_data_parallel_group())
+    assert counts[0].item() == torch.distributed.get_world_size(
+        group=mpu.get_data_parallel_group())

    # Load indexed dataset.
    print_rank_0(' > loading indexed mapping from {}'.format(