From 76e62a5782fc2509ce989fcfc0d0aedc17322b3a Mon Sep 17 00:00:00 2001
From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= <francois@fleuret.org>
Date: Sun, 22 Oct 2023 19:53:59 +0200
Subject: [PATCH] Update.
---
problems.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/problems.py b/problems.py
index 4059856..b8fcdb3 100755
--- a/problems.py
+++ b/problems.py
@@ -313,8 +313,8 @@ class ProblemMixing(Problem):
return y
def start_error(self, x):
- i = torch.arange(self.height).reshape(1,-1,1).expand_as(x)
- j = torch.arange(self.width).reshape(1,1,-1).expand_as(x)
+ i = torch.arange(self.height, device=x.device).reshape(1,-1,1).expand_as(x)
+ j = torch.arange(self.width, device=x.device).reshape(1,1,-1).expand_as(x)
ri = (x == self.height * self.width).long().sum(dim=-1).argmax(-1).view(-1,1,1)
rj = (x == self.height * self.width).long().sum(dim=-2).argmax(-1).view(-1,1,1)
--
2.39.5