Yes, your code should be alright.
However, the .max(1) call will also return the same indices, if you apply it on features directly, so you don’t really need the softmax here.
Also, I would recommend to cast the equality check to float before summing, as it’ll return a torch.uint8 tensor, which might overflow, if you are using large patches or e.g. pixel-wise accuracy:
correct += (predictions == labels.data).float().sum()