Well the only difference I see in training phase VS validation phase is that the batch norm will behave differently in train mode VS eval mode. Could you try running again your code but commenting (just for this test) all the batch norm layers in your model?