Hey @deepak242424
You can use model.apply(function_name) here for batch norm layers
`def set_bn_eval(m):
classname = m.__class__.__name__

if classname.find('BatchNorm') != -1:
    m.eval()`

then call this function by .apply before starting the training.
model.train()
model.apply(set_bn_eval)
Now you can train your model in the same manner as you were. BN layers won’t be trained.
After intermediate validation where you might have used model.eval() then switch to training via model.train() follow that with model.apply(set_bn_eval).
I hope it helps.