Unable to change batch norm mode to training 

How to change batch norm mode to training inside image classification scenario.  I want to update the mean and var for the Pytorch estimator batch norm layers. 


I tried with `estimator.set_batchnorm(train=True)` but it is not working. 