From 54ebb2e5930db8d81a30430190d75a4044874471 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 15 Apr 2019 19:25:36 +0200
Subject: [PATCH] pin_memory=True

---
 test.py  | 2 +-
 train.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/test.py b/test.py
index 222f829a..686a8f5d 100644
--- 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
diff --git a/train.py b/train.py
index 4dca56bc..5bce4adc 100644
--- 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)