How fast does the memory usage increase? Does it keeps increasing after every epoch or stabilized after a few epoches?
It could be due to RRef or distributed autograd context wasn’t deleted in time. It might worth provide an API to block waiting for all RPC workers to clear RRefs and dist autograd contexts. cc @pritamdamania87