- 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)
+ ri = torch.randint(self.height, (nb,)).reshape(nb, 1, 1)
+ rj = torch.randint(self.width, (nb,)).reshape(nb, 1, 1)