Time performance improvement

closes Lightning-AI#12408
isvogor-foi · Mar 23, 2022 · 8396aab · 8396aab
1 parent 1c18d5e
commit 8396aab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_lightning/trainer/supporters.py b/pytorch_lightning/trainer/supporters.py
@@ -73,7 +73,7 @@ def last(self):
     def append(self, x):
         """Add an element to the accumulator."""
         if self.memory is None:
-            self.memory = torch.zeros(self.window_length, *x.shape)
+            self.memory = torch.zeros(self.window_length, *x.shape, x.device)
 
         # ensure same device and type
         if self.memory.device != x.device or self.memory.type() != x.type():