Add non_blocking device transfer in _unpack_minibatch

ContinualAI · Oct 11, 2023 · a09c096 · a09c096
1 parent d7a096e
commit a09c096
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/avalanche/training/supervised/naive_object_detection.py b/avalanche/training/supervised/naive_object_detection.py
@@ -247,10 +247,19 @@ def forward(self):
     def _unpack_minibatch(self):
         # Unpack minibatch mainly takes care of moving tensors to devices.
         # In addition, it will prepare the targets in the proper dict format.
-        images = list(image.to(self.device) for image in self.mbatch[0])
-        targets = [{k: v.to(self.device) for k, v in t.items()} for t in self.mbatch[1]]
-
-        mbatch = [images, targets, torch.as_tensor(self.mbatch[2]).to(self.device)]
+        images = list(
+            image.to(self.device, non_blocking=True) for image in self.mbatch[0]
+        )
+        targets = [
+            {k: v.to(self.device, non_blocking=True) for k, v in t.items()}
+            for t in self.mbatch[1]
+        ]
+
+        mbatch = [
+            images,
+            targets,
+            torch.as_tensor(self.mbatch[2]).to(self.device, non_blocking=True),
+        ]
         self.mbatch = tuple(mbatch)
 
     def backward(self):

diff --git a/avalanche/training/templates/problem_type/supervised_problem.py b/avalanche/training/templates/problem_type/supervised_problem.py
@@ -49,7 +49,7 @@ def _unpack_minibatch(self):
             self.mbatch = mbatch
 
         for i in range(len(mbatch)):
-            mbatch[i] = mbatch[i].to(self.device)  # type: ignore
+            mbatch[i] = mbatch[i].to(self.device, non_blocking=True)  # type: ignore
 
 
 __all__ = ["SupervisedProblem"]