Set drop_last=True for DataLoader (#1889)

aws · Dec 30, 2020 · c6cf46d · c6cf46d
1 parent c608af2
commit c6cf46d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/...g/distributed_training/pytorch/model_parallel/bert/bert_example/sagemaker_smp_pretrain.py b/...g/distributed_training/pytorch/model_parallel/bert/bert_example/sagemaker_smp_pretrain.py
@@ -92,7 +92,7 @@ def create_pretraining_dataset(input_file, max_pred_length, shared_list, args, w
     train_dataloader = DataLoader(train_data, sampler=train_sampler,
                                   batch_size=args.train_batch_size * args.n_gpu,
                                   num_workers=4, worker_init_fn=worker_init,
-                                  pin_memory=True)
+                                  pin_memory=True, drop_last=True)
     return train_dataloader, input_file
 
 
@@ -696,7 +696,7 @@ def main():
             train_dataloader = DataLoader(train_data, sampler=train_sampler,
                                         batch_size=args.train_batch_size * args.n_gpu,
                                         num_workers=4, worker_init_fn=worker_init,
-                                        pin_memory=True)
+                                        pin_memory=True, drop_last=True)
             # shared_file_list["0"] = (train_dataloader, data_file)
         else:
             train_dataloader = restored_data_loader