[fix-daily] Use composer get_model_state_dict instead of torch's

mosaicml · Jul 25, 2024 · 55f0b7d · 55f0b7d
1 parent c5542a3
commit 55f0b7d
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/checkpoint/test_download.py b/tests/checkpoint/test_download.py
@@ -8,7 +8,7 @@
 import pytest
 import torch
 
-from composer.checkpoint import download_monolithic_checkpoint
+from composer.checkpoint import download_monolithic_checkpoint, get_model_state_dict
 from composer.utils import dist
 from tests.checkpoint.helpers import init_model
 from tests.common.markers import world_size
@@ -26,8 +26,7 @@ def test_download_monolithic_checkpoint(world_size: int, rank_zero_only: bool):
         use_fsdp = True
     fsdp_model, _ = init_model(use_fsdp=use_fsdp)
 
-    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
-    state = get_model_state_dict(fsdp_model, options=StateDictOptions(full_state_dict=True))
+    state = get_model_state_dict(fsdp_model, sharded_state_dict=False)
 
     checkpoint_filename = 'state_dict'
     save_filename = os.path.join(tmp_dir.name, checkpoint_filename)