


loss = x.sum()
      >>> loss.backward()
      >>> optim.step()



Warning:
 The optimizer must be initialized *after* the module has been
 wrapped, since FSDP will shard parameters in-place and this will
 break any previously initialized optimizers.

Warning:
 If the destination CUDA device has ID "dev_id", either (1)
 "module" should already be placed on that device, (2) the device
 should be set using "torch.cuda.set_device(dev_id)", or (3)
 "dev_id" should be passed into the "device_id" constructor
 argument. This FSDP instance's compute device will be that
 destination device. For (1) and (3), the FSDP initialization
 always occurs on GPU. For (2), the FSDP initialization happens on
 "module" 's current device, which may be CPU.

Warning:
 FSDP currently does not support gradient accumulation outside
 "no_sync()" when using CPU offloading. Trying to do so yields
 incorrect results since FSDP will use the newly-reduced gradient
