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?