Hi @xieydd, there have a been two bug fixes for these issues, both landed in master.  One of the fixes updated device affinity for enabling distributed mode, and another updated device affinity for converting a model to QAT.  Would you mind trying with the nightly build?  All of the fixes will be a part of the upcoming v1.6 release as well.
If it’s hard to get the nightly, for a quick fix, inserting model.to(device) after line 71 should work.