pin_memory=True
This commit is contained in:
parent
e3f0b0248c
commit
54ebb2e593
2
test.py
2
test.py
|
@ -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
|
||||||
|
|
2
train.py
2
train.py
|
@ -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)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue