commit
b22ee4dd5f
|
@ -65,6 +65,7 @@ def build_dataloader(config, mode, device, logger, seed=None):
|
|||
loader_config = config[mode]['loader']
|
||||
batch_size = loader_config['batch_size_per_card']
|
||||
drop_last = loader_config['drop_last']
|
||||
shuffle = loader_config['shuffle']
|
||||
num_workers = loader_config['num_workers']
|
||||
if 'use_shared_memory' in loader_config.keys():
|
||||
use_shared_memory = loader_config['use_shared_memory']
|
||||
|
@ -75,14 +76,14 @@ def build_dataloader(config, mode, device, logger, seed=None):
|
|||
batch_sampler = DistributedBatchSampler(
|
||||
dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
shuffle=shuffle,
|
||||
drop_last=drop_last)
|
||||
else:
|
||||
#Distribute data to single card
|
||||
batch_sampler = BatchSampler(
|
||||
dataset=dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False,
|
||||
shuffle=shuffle,
|
||||
drop_last=drop_last)
|
||||
|
||||
data_loader = DataLoader(
|
||||
|
|
Loading…
Reference in New Issue