Thanks @chetan_patil. As for my use case right now, I didn’t need batch normalization. Since my architecture is absolutely same, I think I should expect  torch.save() to do the job i.e. using the saved model for another iteration for more number of samples.