- 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)