From 0b2ddd4277f2d3ced5d52e67e1bb1762da223ba9 Mon Sep 17 00:00:00 2001 From: outerskyb Date: Mon, 8 Feb 2021 01:20:51 +0900 Subject: [PATCH 1/2] fix multiple devices problem --- models/morn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/models/morn.py b/models/morn.py index 8af90ce..706ac19 100644 --- 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) From f7e616fc51f3c9c14af3205a918a8275f007a6d4 Mon Sep 17 00:00:00 2001 From: outerskyb Date: Mon, 8 Feb 2021 11:04:32 +0900 Subject: [PATCH 2/2] fix wrong spaces --- models/morn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/morn.py b/models/morn.py index 706ac19..b360eac 100644 --- a/models/morn.py +++ b/models/morn.py @@ -110,7 +110,7 @@ def forward(self, x, test, enhance=1, debug=False): v_max = offsets_max.data[i] v_min = offsets_min.data[i] - if self.cuda: + 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))