kazem:

checkpoint = torch.load(args.load_path) # , map_location=map_location)


This line might cause a problem if the model was saved from a device that is not available on the machine that loads the model. But it should be OK in your case, as the model was saved from rank 0 (i.e., “cuda:0”), whose device is available in both envs. However, without map_location, it means the two DDP processes in testing are operating on the same GPU? That could also cause problems.