From 40a932f1373ac42654c85bf4fe2924ba5a5206a3 Mon Sep 17 00:00:00 2001 From: davidtweedle Date: Thu, 5 Dec 2024 14:18:22 -0400 Subject: [PATCH] Typo in FSDP definition for ogbg --- algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py b/algorithmic_efficiency/workloads/ogbg/ogbg_pytorch/workload.py index 1f54b75fd..f3af1edc0 100644 --- 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