One approach would be to use pass a few validation batches in train mode to adjust the running estimates.
I’m not completely sure, if that would be some kind of data leakage, but you might give it a try and see if it helps.