diff --git a/avalanche/training/supervised/naive_object_detection.py b/avalanche/training/supervised/naive_object_detection.py index 3c13d8e17..e0c60bede 100644 --- 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 index fff259ee4..b2a264b63 100644 --- 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"]