Hello,
I’ve been struggling for a while with the runtime error. I have visited most, if not all, of the posts on this forum that has this error, and tried to replicate the solution, but I still cannot fix it. Any suggestions would be appreciated! 
Essentially, I have a custom PyTorch dataset class, where I set requires_grad=False for both input x and groundtruth y. I read that this is a must in order to iterate through the dataloader, and this seems to work.
But I also read on this forum that I have to set input x with requires_grad=True before passing into the network, since I had previously set it to False.
But as for groundtruth, despite we use it in the loss function, do we still need to set requires_grad=True for that? I’ve tried that and still get the same runtime error.
Essentially, my code runs past the validation sanity check. It was able to generate the losses twice at Epoch 0, before running into an error.
For all the losses generated at Epoch 0, I see that there’s grad_fn being recorded. So I really don’t understand where there is an error.
Epoch 0:   0%|                                                                                                     | 0/16 [00:00<00:00, 5011.12it/s]

{'loss1': tensor(1.0887e+15), 'loss2': tensor(11.4520, grad_fn=<DivBackward0>), 'loss3': tensor(2.0911, grad_fn=<LogBackward0>)}
tensor(1.0887e+15, grad_fn=<AddBackward0>)

{'loss1': tensor(1.0894e+15), 'loss2': tensor(11.4538, grad_fn=<DivBackward0>), 'loss3': tensor(2.0907, grad_fn=<LogBackward0>)}
tensor(1.0894e+15, grad_fn=<AddBackward0>)

Traceback (most recent call last):
....

Below is the training step:
    def training_step(self, train_batch, batch_idx):
        
        # x = (2,8,img_yaxis,img_xaxis) images, y = groundtruth
        x, y = train_batch

        # Set requires_grad = True
        x = torch.tensor(x, requires_grad=True)
        
        out = self.network(x)

        loss_func = ls.Loss(y[0])
    
        # Out = Tensor([probability_map, gene_map])
        # probability map: 3D (3 batches, 572, 572)
        # barcode gene map: 4D (3 batches, 16, 572, 572)
        probability_map = out[:, 0, :,:]
        gene_map = out[:, 1:17, :,:]
        
        # Now probability_map is 3D (3 batches, 572, 572).
        # So we cast it back to 4D (3 batches, 1, 572, 572).
        probability_map = probability_map[:, None, :, :]
        
        num_batches = probability_map.shape[0]
        
        # FAST COMPUTATION OF LOSS WITHOUT FOR-LOOP using MAP()
        loss = list(map(loss_func, probability_map, gene_map))

        for idx in range(num_batches):
            
            self.log('train_loss', loss[idx])
        
        return torch.mean(torch.FloatTensor(loss))