Update thunder/distributed/transforms/fsdp_v2.py

Lightning-AI · Aug 19, 2024 · 2357aad · 2357aad
1 parent f05768e
commit 2357aad
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thunder/distributed/transforms/fsdp_v2.py b/thunder/distributed/transforms/fsdp_v2.py
@@ -223,7 +223,7 @@ def transform_module(
                     if p_orig.device.type != "meta":
                         p_meta = torch.nn.Parameter(p.to(device="meta"), requires_grad=p.requires_grad)
                         p_meta._thunder_device = p_orig.device
-                        setattr(submodule, base_pn, p_meta)
+                        submodule.register_parameter(base_pn, p_meta)
                     else:
                         p_orig._thunder_device = self.device