Since you mentioned federated learning, shouldn’t the data transfer be in a distributed environment? How about the following code snippet?
    # Flattening all the parameters of the cloud model into a contiguous buffer to prepare for data transfer.
    flat_params = torch.cat([p.data.view(-1) for p in model.parameters()])

    # broadcast the tensors or call process group send/recv?
    ...
    
    # Copy the parameters to the client model layer by layer.
    offset = 0
    for p in module.parameters():
        p.data = flat_params[offset : offset + p.numel()].view_as(p)
        offset += p.numel()