pin_memory=True

2019-04-15 19:25:36 +02:00 · 2019-04-15 19:25:36 +02:00 · 54ebb2e593
parent e3f0b0248c
commit 54ebb2e593
2 changed files with 2 additions and 2 deletions
--- a/test.py
+++ b/test.py
@ -48,7 +48,7 @@ def test(
    dataloader = DataLoader(dataset,
                            batch_size=batch_size,
                            num_workers=4,
-                            pin_memory=False,
+                            pin_memory=True,
                            collate_fn=dataset.collate_fn)

    seen = 0
--- a/train.py
+++ b/train.py
@ -93,7 +93,7 @@ def train(
                            batch_size=batch_size,
                            num_workers=num_workers,
                            shuffle=False,
-                            pin_memory=False,
+                            pin_memory=True,
                            collate_fn=dataset.collate_fn,
                            sampler=sampler)