I’m not sure, how your CustomLoader2 works, but if it’s some kind of Dataset, it should be alright.
Some small notes:

usually you don’t need to set shuffle=True for your test DataLoader, since the accuracy and loss won’t depend on the order of your samples. It won’t break anything, but you could just remove it.
you forgot the parentheses in model.train(), so currently your model might get stuck in model.eval(), which will disable the dropout layers
it’s recommended to call the model directly instead of the forward method, so that all hooks will be properly registered. Probably not important in your current script, but anyway I would stick to it.
if I’m not mistaken you don’t need to call torch.exp(log_ps) to get the predictions, since your model outputs raw logits.