def METHOD_NAME(data):
    # with FSDP, module params will be on CUDA, so we cast them back to CPU
    # so that the RNG is consistent with and without FSDP
    data.copy_(data.cpu().METHOD_NAME(mean=0.0, std=0.02).to(data.device))