From 528e690f39966974f3e17509c4397002ee9db45a Mon Sep 17 00:00:00 2001
From: janEbert <janpublicebert@posteo.net>
Date: Tue, 3 Jun 2025 11:59:13 +0200
Subject: [PATCH] Use meta device initialization
---
pytorch-fsdp-example/main.py | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/pytorch-fsdp-example/main.py b/pytorch-fsdp-example/main.py
index 6d6eeb3..cc845ed 100644
--- a/pytorch-fsdp-example/main.py
+++ b/pytorch-fsdp-example/main.py
@@ -375,8 +375,12 @@ def main():
train_dset, valid_dset, test_dset = prepare_datasets(args, device)
- model = build_model(args)
+ 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
--
GitLab