Generate 2 D data

Hi, i want to generate the u(N,2) data to feed the NN and also want to save the rx and rt as x_space and t_time to again feed in to get the differential of u wrt to rt and rx.

rx = torch.tensor((np.arange(0,1,.01).reshape(100,1)), requires_grad=True, dtype=torch.float32)
rt = torch.tensor((np.arange(0,30,.3).reshape(100,1)), requires_grad=True, dtype = torch.float32)

def differential(u_pred,x,t):

    z = torch.ones(u_pred.shape).requires_grad_(True)
    dx = torch.autograd.grad(u_pred, x, grad_outputs=z, create_graph=True)[0]
    dt = torch.autograd.grad(u_pred, t, grad_outputs=z, create_graph=True)[0]
    dxx = torch.autograd.grad(dx, z, grad_outputs=torch.ones(dx.shape), create_graph=True)[0]
    f = dxx-dt
    f = pn.MSE(f , torch.zeros(len(f)))
    
    return f

Is it a better way to do this so that i can save u, t and x to feed in?