Henry_Chibueze:

I tried print the pred before and after the _, pred = torch.max(pred, 1)
before it, the pred seems to always have negative values at index 0 and positive values at index 1 per batch and it’s not a deterministic behavior coz it happens vise versa.
These values repeat themselves a lot and they barely change
after it, the pred is only one class per batch (either 0s or 1s)


Before you do _, pred = torch.max(pred, 1),  pred  is an Nx2 array where, within any one batch, either the 0th column is all negative and the 1st column is all positive, or vice versa.
When you do  _, pred = torch.max(pred, 1) you update the value of pred to be an Nx1 array with the following property: pred[i] is the index of the column which had the maximum value in row i of the original pred. Since the original pred is as you described above, obviously this results in a new pred which is all 0s (if the original pred had positive values at index 0) or all 1s (if the original pred had positive values at index 1). See the documentation for torch.max. Note that there are two variants of torch.max described in that page; your usage is the second one.
If this processing of pred is as you expect it to be, you have to investigate why the original pred has the same sign for each column, within any one batch. If this processing of pred is not as per your expectation (but the original value of pred is fine), then you may want to check the other parameters of torch.max to see if you can use them to get what you want.