Akshay_A_J:

I am not understanding why there is problem with using torch.max?


You can calculate the gradients for the returned values in torch.max, but not the returned indices as this operation is not differentiable:
x = torch.randn(10, 10, requires_grad=True)
val, idx = torch.max(x, dim=1)
print(val, idx)
> tensor([1.4675, 2.2521, 1.3766, 1.8801, 1.5437, 1.3367, 0.8476, 0.7166, 1.0904,
        1.5543], grad_fn=<MaxBackward0>) tensor([3, 9, 4, 8, 3, 1, 0, 2, 1, 2])

As you can see, the val tensor contains a valid .grad_fn, while the idx tensor does not.
Also, where did you use the .data attribute? It shouldn’t be used at all generally, as Autograd cannot track the operations on this attribute, but I also thought you’ve already removed the usage in the previous code snippet.