The code might has memory leak

[ray8828]

Hi, thanks for sharing this work, seems the training can't stand for several epoches, it can cause oom on a 48GB RAM machine even, seems the code would result in host memory leak

[pufferandrr]

Have you solved this problem? I met the same situation with you