pin_memory=True

This commit is contained in:
Glenn Jocher 2019-04-15 19:25:36 +02:00
parent e3f0b0248c
commit 54ebb2e593
2 changed files with 2 additions and 2 deletions

View File

@ -48,7 +48,7 @@ def test(
dataloader = DataLoader(dataset, dataloader = DataLoader(dataset,
batch_size=batch_size, batch_size=batch_size,
num_workers=4, num_workers=4,
pin_memory=False, pin_memory=True,
collate_fn=dataset.collate_fn) collate_fn=dataset.collate_fn)
seen = 0 seen = 0

View File

@ -93,7 +93,7 @@ def train(
batch_size=batch_size, batch_size=batch_size,
num_workers=num_workers, num_workers=num_workers,
shuffle=False, shuffle=False,
pin_memory=False, pin_memory=True,
collate_fn=dataset.collate_fn, collate_fn=dataset.collate_fn,
sampler=sampler) sampler=sampler)