model.train(True)

for m in model.modules():
  if isinstance(m, nn.BatchNorm1d) or isinstance(m, nn.BatchNorm2d):
    m.eval()

# Run your training here

You can further set requires_grad=False for your BN layers, but that only effects weights + biases.