Thanks a lot ptrblk. I eliminate learning rate scheduler and the problem goes away! So I did not reproduced the situation to look for overheating issue. I’m wondering if I do something wrong with learning rate scheduler?
Currently I replaced scheduler.step() with:
if (epoch + 1) % 5000 == 0:
    for g in optimizer.param_groups:
        g['lr'] = g['lr'] * 0.75

and I’m waiting to see whether the problem comes in or not, but tests are a little time consuming…