I tried your method after a few minor correction but it still gives me the same erroneous result. I use this resnet script to call the model.

You mean you saw error by running the script as is? What error did you see and what fix did you applied?

Also I only want to save and load the state_dict and not the entire model since it takes a lot of space.

It should be doable by just modifying two lines (save and load).

When I train DDP with 8 gpus and test DDP with 8 gpus later, there is no issue. However, when I train DDP with 8 gpus and test DDP with 2 gpus later the problem occurs.

The resnet link you posted points to torchvision resnet, so the model only lives on a single device. How did you go from training on 8 gpus to testing on 2 gpus? Did you do the following?

After training, use only rank 0 to save ddp.module to file.
For testing, as you no longer need comm across models, you don’t need DDP. You can spawn two processes, each load the saved module from file to its dedicated device by setting map_reduce. And use sth like all_gather to collect loss/accuracy data to rank 0?