@kunasiramesh, @Gkv The memory issue might be related to the training procedure or another part of the code.
Could you post the code so that we can have a look?
Usually the computation graph is unintentionally stored somewhere, e.g. by using losses += loss instead of losses += loss.item().