We found the reason of this slight increase. The state variable used for saving the checkpoint as above was not nullified during training epochs.
Adding state = None after saving checkpoint helped.
Issue observed only with models with batchnorm and potentially has to do with the buffers as, the same memory increase is not observed with track_running_stats=False.