correct test_distribtued

pytorch · Dec 10, 2024 · 8aae3d9 · 8aae3d9
1 parent 22f3f17
commit 8aae3d9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/torchtune/training/test_distributed.py b/tests/torchtune/training/test_distributed.py
@@ -243,6 +243,7 @@ def test_lora_state_dict(self):
             fsdp_model_to_load.parameters(), weight_decay=0.01, lr=0.01
         )
         training.load_from_full_optimizer_state_dict(
+            fsdp_model_to_load,
             fsdp_optim_to_load,
             # mimic mmap=True where every rank see full SD
             copy.deepcopy(self._broadcast_full_state_dict(optim_full_sd)),