Hey @john, is any part of this model using DataParallel, DistributedDataParallel, or torch.distributed.rpc? Any reason for tagging this question with “distributed-rpc”?
The format of the code looks distorted, and will be hard to debug. Could you please share a properly-formatted self-contained example?