The loss calculation for nn.BCELoss looks wrong, as this criterion expects the model outputs to be probabilities provided via a sigmoid activation, while you are applying torch.max on it.
Besides that the code looks alright and I cannot find anything obviously wrong.
I think you’ll find some posts here in the forum which are discussing this particular dataset, which might contain some code as a starter.