Hello there !
I see you’re using an MLP, depending on the activation functions you have, if you use 2nd order derivatives for ReLU activations, you should have all zeros, be careful with that. Tanh or sigmoid activation should do the trick.
In addition, if you want to differentiate your neural net wrt t or x, I advise you to be a bit more careful with how you write your derivatives.
Here’s a little example:
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.l1=nn.Linear(2,80)
        self.mi=nn.Linear(80,80)
        self.mi1=nn.Linear(80,80)
        self.mi2=nn.Linear(80,40)
        self.ol=nn.Linear(40,1)
        self.th=nn.Tanh()
            
    def forward(self,x,t):
        u=torch.cat((x,t),1)
        u=self.th(self.l1(u))
        u=self.th(self.mi(u))
        u=self.th(self.mi1(u))
        u=self.th(self.mi2(u))
        u=self.th(self.ol(u))
        
        return u
def f(x,t):
    u = net(x,t)
    u_t = grad(u, t, 
                     create_graph=True,
                     grad_outputs=torch.ones_like(u), 
                     allow_unused=True)[0]
    u_x = grad(u, x, 
                     create_graph=True,
                     grad_outputs=torch.ones_like(u), 
                     allow_unused=True)[0]
    u_xx = grad(u_x, x, 
                     create_graph=True,
                     grad_outputs=torch.ones_like(u), 
                     allow_unused=True)[0]
    w = torch.tensor(0.01/np.pi)
    f = u_t + u*u_x - w*u_xx
    return f

Here’s an example for Burgers’ equation. I got the original code from NNets-and-Diffeqns/burgernet.py at master · dbgannon/NNets-and-Diffeqns · GitHub 44
It’s explained here: Notes on Deep Learning and Differential Equations. – Cloud Computing For Science and Engineering 23
Hope it helps !