Use meta device initialization

pytorch-fsdp-example/main.py

@@ -375,8 +375,12 @@ def main():
 
     train_dset, valid_dset, test_dset = prepare_datasets(args, device)
 
+    with torch.device('meta'):
         model = build_model(args)
     model = distribute_model(model)
+    # Put model from meta device to actual device.
+    model.to_empty(device=device)
+    model.reset_parameters()
     loss_func = torch.nn.CrossEntropyLoss()
 
     lr = args.lr