pytorch测试内存泄漏

import psutil
    process = psutil.Process(os.getpid())
    for _ in range(10):
        for i, batch in tqdm(enumerate(train_data)):
            img, text = batch
            mm_info = process.memory_full_info()
            print(i, mm_info.uss / 1024 / 1024, "MB")