Burgers 2D Notebook not running properly

[lcdbezerra]

When executing Burger 2D.ipynb I get the following error:

model.train()
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [5], in <cell line: 1>()
----> 1 model.train()

File ~/code/pinn-pytorch/grad/pytorchPDE.py:111, in PytorchPDE.train(self, epochs, save_step, update_step, upload)
    107 self.optimizer.zero_grad()
    109 # forward + backward + optimize
--> 111 loss = self.loss()
    113 # if(upload):
    114 #     neptune.log_metric('loss', loss.item())
    116 loss.backward()

File ~/code/pinn-pytorch/grad/pytorchPDE.py:56, in PytorchPDE.loss(self)
     48             second_grads = None
     50 #         print('train_tag_min: ', self.train_tag[:,-1].min())
     51 #         print('train_tag_max: ', self.train_tag[:,-1].max())
     52 #         print('predicts_min: ', predicts.min())
     53 #         print('first_grads_min: ', first_grads.min())
     54 #         print('second_grads_min: ', second_grads[0].min(), second_grads[0].min())
---> 56         domain_loss = self.pde(train_points, predicts, first_grads, second_grads)*self.train_tag[:,-1]
     57         boundary_loss = 0
     58         for i, bc in enumerate(self.data.bcs):

Input In [4], in pde(x, y, first_grads, second_grads)
      1 def pde(x,y,first_grads,second_grads):
      2 #     print(second_grads)
----> 3     return first_grads[:,0] + y[:,0]*first_grads[:,1] - (0.01/np.pi)*second_grads[1][:,1]

IndexError: tuple index out of range

Package versions:

python=3.9.12
torch=1.11.0