I am train model with distributed_training_mnist.py and curious if I am can use somethin like NormalSampler with it?