


...
       >>> fsdp_model = FSDP(module, param_init_fn=my_init_fn, auto_wrap_policy=size_based_auto_wrap_policy)
       >>> print(next(fsdp_model.parameters()).device) # current CUDA device
       >>> # With torchdistX
       >>> module = deferred_init.deferred_init(MyModule, device="cuda")
       >>> # Will initialize via deferred_init.materialize_module().
       >>> fsdp_model = FSDP(module, auto_wrap_policy=size_based_auto_wrap_policy)




  * **device_id** (*Optional**[**Union**[**int**,
    **torch.device**]**]*) -- An "int" or "torch.device"
    describing the CUDA device the FSDP module should be moved to
    determining where initialization such as sharding takes place.
    If this argument is not specified and "module" is on CPU, we
    issue a warning mentioning that this argument can be specified
    for faster initialization. If specified, resulting FSDP
    instances will reside on this device, including moving ignored
