I don’t know if this is from posting here, but you would want the model.eval() to happen outside the training loop, right? Also, you would want the validation to happen on the validation loop, probably.
torch.no_grad() saves memory and, sometimes, time, but the recorded autograd graphs should be deallocated when the validation leaves loss_for_batch because all tensors (Y_prob and loss in particular) go out of scope regardless of whether you have done the backward, so at the end of each iteration you will be in the same situation with or without no_grad in the loss_for_batch function. There used to be bugs around this a longish time ago, but I think the above would work.
That said, I would probably add the no_grad context (or a torch.set_grad_enabled(torch.is_grad_enabled() and optimizer is not None) context if you want to handle with gradient and without in one bit of code).
Best regards
Thomas