Skip to content

Commit

Permalink
Drop last in distributed mode
Browse files Browse the repository at this point in the history
  • Loading branch information
sarlinpe committed Jan 10, 2024
1 parent 4a82835 commit 2079d6c
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions gluefactory/datasets/base_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,9 +161,12 @@ def get_data_loader(self, split, shuffle=None, pinned=False, distributed=False):
except omegaconf.MissingMandatoryValue:
batch_size = self.conf.batch_size
num_workers = self.conf.get("num_workers", batch_size)
drop_last = True if split == "train" else False
if distributed:
shuffle = False
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
sampler = torch.utils.data.distributed.DistributedSampler(
dataset, drop_last=drop_last
)
else:
sampler = None
if shuffle is None:
Expand All @@ -178,7 +181,7 @@ def get_data_loader(self, split, shuffle=None, pinned=False, distributed=False):
num_workers=num_workers,
worker_init_fn=worker_init_fn,
prefetch_factor=self.conf.prefetch_factor,
drop_last=True if split == "train" else False,
drop_last=drop_last,
)

def get_overfit_loader(self, split):
Expand Down

0 comments on commit 2079d6c

Please sign in to comment.