Merge pull request #145 from Outerskyb/master

Fix multiple devices problems
Canjie-Luo · Dec 6, 2022 · 410596f · 410596f
2 parents 2cd40c4 + f7e616f
commit 410596f
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/morn.py b/models/morn.py
@@ -110,7 +110,10 @@ def forward(self, x, test, enhance=1, debug=False):
 
                 v_max = offsets_max.data[i]
                 v_min = offsets_min.data[i]
-                img_offsets = (offsets_grid[i]).view(1, self.targetH, self.targetW).data.cpu().add_(-v_min).mul_(1./(v_max-v_min))
+                if self.cuda:
+                    img_offsets = (offsets_grid[i]).view(1, self.targetH, self.targetW).data.cuda().add_(-v_min).mul_(1./(v_max-v_min))
+                else:
+                    img_offsets = (offsets_grid[i]).view(1, self.targetH, self.targetW).data.cpu().add_(-v_min).mul_(1./(v_max-v_min))
                 img_offsets = to_pil_image(img_offsets)
                 img_offsets = np.array(img_offsets)
                 color_map = np.empty([self.targetH, self.targetW, 3], dtype=int)