reward = 3
def custom_loss(outputs,target):
        outputs  = F.softmax(outputs,dim=1)
        outputs, reservation = outputs[:,:-1], outputs[:,-1]
        gain = torch.gather(outputs, dim=1, index=target.unsqueeze(1)).squeeze()
        return loss

I am trying to implement custom loss but it was giving runtime cuda error so i remove lines step by step and tried to run CNN. It works fine. BUT
The moment i enter
        gain = torch.gather(outputs, dim=1, index=target.unsqueeze(1)).squeeze()
This line it through
RuntimeError: CUDA error: device-side assert triggered