the indices returned by torch.argmax or by torch.max as the second return value are not differentiable as seen here:
outputs = torch.randn(10, 10, requires_grad=True)
_, preds = torch.max(outputs, dim=1)
print(preds.grad_fn)
> None

To calculate the dice loss you could pass F.softmax(outputs, dim=1) as the predictions.
Also, don’t use the .data attribute, as it might have unwanted side effects since Autograd won’t be able to track these operations.