I am working on creating custom loss function in pytorch, below the code:
class CustomSoftmaxLoss(torch.nn.Module):
    def __init__(self):
        super(CustomSoftmaxLoss, self).__init__()

    def forward(self, predictions, targets):
        sample_size = predictions.size()[0]

        # Fix numerical instability
        predictions -= predictions.max(dim=1)[0].view(-1, 1)

        # Run predicitions through softmax
        softmax = torch.exp(predictions.float())
        softmax = softmax / softmax.sum(dim=1).view(-1, 1)

        # Cross-entropy loss
        loss = -torch.log(softmax[torch.arange(sample_size).type(torch.LongTensor), targets]).sum()
        loss /= sample_size

        return loss

# Test
predictions = torch.from_numpy(np.array([[1,2,3], [3,15,7], [10,4,2], [2,3,4]]))
targets = torch.from_numpy(np.array([0,1,2,2]))

predictions = Variable(predictions)
targets = Variable(targets)                          

criterion = CustomSoftmaxLoss()
a         = criterion(predictions, targets)
print(a)

The code above is working. However, if I enable CUDA like below, there is error
criterion = CustomSoftmaxLoss().cuda()
a         = criterion(predictions.cuda(), targets.cuda())

TypeError: Performing basic indexing on a tensor and encountered an error indexing dim 0 with an object of type torch.LongTensor. The only supported types are integers, slices, numpy scalars, or if indexing with a torch.cuda.LongTensor or torch.cuda.ByteTensor only a single Tensor may be passed.
How to fix this please?
Thanks.