I have a tensor of shape (50,50,2). That is filled with values of a grid in (x,y). Like this.
r =0.5
#Sample 3D space
min = -10
max = 10
complexnum = 50j

#Sample 3D space
X,Y  = np.mgrid[min:max:complexnum,min:max:complexnum,] * r

T = np.sin(3*X) * np.sin(3*Y)


XY = np.concatenate((X.reshape((X.shape[0],X.shape[0],1)),Y.reshape((X.shape[0],X.shape[0],1))),2)
XY = torch.from_numpy(XY.astype(np.float32)).to(device)
T = T.reshape((T.shape[0],T.shape[0],1))
T = torch.from_numpy(T.astype(np.float32)).to(device)

I created a Neural Network that can take the XY as input and T as an output. In that is creating a function f(x,y). Now I need to create a way to train that NN using batches. The total amount of samples should be 2500 = 50 x 50. I want to divide the XY and T tensor into batches. I have tried DataLoader and creating a data set class with no luck. Is there a way I can achieve this? Are my Tensor shapes correct?
Sorry for this very basic question, I have been struggling to work with higher dimensions.
Are the questions badly explained? I would have guessed someone would reply by now.