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.