Hey @jinfagang
Are you asking rank 0 to save the model and then all ranks to load from that checkpoint? If so, you might need to provide a map_location arg when calling torch.load, otherwise, it might load to the device where it was saved. If this still doesn’t fix the problem, I would try first move the model to CPU, and then save it. When loading, always load it to CPU on all ranks, and then explicitly move it to the destination device.