Typo in FSDP definition for ogbg

mlcommons · Dec 5, 2024 · 40a932f · 40a932f
1 parent dbfd233
commit 40a932f
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py b/algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py
@@ -167,12 +167,11 @@ def init_model_fn(
                 size_based_auto_wrap_policy, min_num_params=2 ** 10
                 )
         model = FSDP(
-                self._model,
+                model,
                 use_orig_params=True,
                 auto_wrap_policy=auto_wrap_policy,
                 device_id=RANK
                 )
-
       else:
         model = torch.nn.DataParallel(model)
     return model, None