Solved by ptrblck in post #6 
              
              
                @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().