I also tried a “raw” implementation (without the use of nn.Module) as it is in the original Tensorflow implementation. However I obtain the same result:
# this code is the continuation of the previous one
# PyTorch raw:
#-------------

def initialize_NN_torch(weights_tf):        
    weights_torch = []
    biases_torch = []
    
    for w_tf in weights_tf:
        w_torch = torch.from_numpy(
            w_tf
        ).float()
        w_torch.requires_grad = True
        b_torch = torch.zeros(w_tf.shape[1])
        b_torch.requires_grad = True

        weights_torch.append(w_torch)
        biases_torch.append(b_torch)  
        
    return weights_torch, biases_torch

def neural_net_torch(H, weights, biases):
    num_layers = len(weights) + 1

    for l in range(0,num_layers-2):
        W = weights[l]
        b = biases[l]
        H = torch.tanh(torch.add(torch.mm(H, W), b))
        
    W = weights[-1]
    b = biases[-1]
    Y = torch.add(torch.mm(H, W), b)
    
    return Y

# data
# convert inputs to torch
x_t2 = torch.from_numpy(x_train).float()
y_t2 = torch.from_numpy(y_train).float()
t_t2 = torch.from_numpy(t_train).float()
x_t2.requires_grad = True
y_t2.requires_grad = True
t_t2.requires_grad = True

# preprocess the data
X2 = torch.cat([x_t2, y_t2, t_t2], dim=1)
lb2, _ = X2.min(0)
ub2, _ = X2.max(0)
H2 = 2.0*(X2 - lb2)/(ub2 - lb2) - 1.0

weights_torch, biases_torch = initialize_NN_torch(weights)

psi_and_p = neural_net(H2, weights_torch, biases_torch)
psi = psi_and_p[:,0:1]
p = psi_and_p[:,1:2]

# create a dummy tensor for computing the gradient
dummy = torch.ones((5,1), requires_grad=False)

# gradients
u = torch.autograd.grad(psi, y_t2, grad_outputs=dummy, create_graph=True, only_inputs=True)[0]
v = - torch.autograd.grad(psi, x_t2, grad_outputs=dummy, create_graph=True, only_inputs=True)[0]

print("Psi:", psi)
print("U:", u)
print("V:", v)

For sake of completeness, I attach the results from the three approaches:
Tensorflow:
Psi: [[ 0.05774455]
 [ 0.08581608]
 [-0.13483115]
 [-0.07894165]
 [-0.00312972]]
U: [[0.048948  ]
 [0.05040096]
 [0.03029619]
 [0.04972227]
 [0.04007644]]
V: [[-0.07884889]
 [-0.07497925]
 [-0.07125001]
 [-0.07668343]
 [-0.0951325 ]]

PyTorch:
Psi: tensor([[ 0.0577],
        [ 0.0858],
        [-0.1348],
        [-0.0789],
        [-0.0031]], grad_fn=<SliceBackward>)
U: tensor([[ 0.0489],
        [-0.0578],
        [ 0.0303],
        [-0.0616],
        [ 0.0401]], grad_fn=<SliceBackward>)
V: tensor([[-0.0788],
        [ 0.1062],
        [ 0.1445],
        [-0.0767],
        [-0.0951]], grad_fn=<NegBackward>)


PyTorch “raw”:
Psi: tensor([[ 0.0577],
        [ 0.0858],
        [-0.1348],
        [-0.0789],
        [-0.0031]], grad_fn=<SliceBackward>)
U: tensor([[ 0.0489],
        [-0.0578],
        [ 0.0303],
        [-0.0616],
        [ 0.0401]], grad_fn=<SliceBackward>)
V: tensor([[-0.0788],
        [ 0.1062],
        [ 0.1445],
        [-0.0767],
        [-0.0951]], grad_fn=<NegBackward>)