- i = torch.arange(self.height).reshape(1,-1,1).expand(nb,self.height,self.width)
- j = torch.arange(self.width).reshape(1,1,-1).expand(nb,self.height,self.width)
+ i = (
+ torch.arange(self.height)
+ .reshape(1, -1, 1)
+ .expand(nb, self.height, self.width)
+ )
+ j = (
+ torch.arange(self.width)
+ .reshape(1, 1, -1)
+ .expand(nb, self.height, self.width)
+ )