checkpoint = torch.load(CHECKPOINT, map_location=map_location)
   ddp_model.load_state_dict(checkpoint['state_dict'])
   powersgd_hook = checkpoint['comm_hook']
   powersgd_state = checkpoint['comm_hook_state']

   ddp_model.register_comm_hook(powersgd_state, powersgd_hook)

   if rank == 0:
       os.remove(CHECKPOINT)

   cleanup()

if name == "main":
       n_gpus = torch.cuda.device_count()
       assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
       world_size = n_gpus
       run_demo(demo_serialization, world_size)
Acknowledgements
Many thanks to PowerSGD paper author Thijs Vogels for the code
review on PowerSGD communication hook, as well as the comparison
experiments, which show that the performance of PowerSGD communication
hook is on par with the implementation in the original paper.