I think it’s quite good.
The only comments I would make are:
Do total_loss += loss.item() to convert the loss to a python number directly. This will make sure you don’t build the autograd graph for things that don’t need it.
You can wrap you validate function with @torch.no_grad() to disable the autograd in the whole function if you want (you already do it for the model which is the most important part).
Looks good otherwise.