fix assert error

AlibabaPAI · Sep 24, 2024 · c764c97 · c764c97
1 parent b866dde
commit c764c97
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchacc/dist/fsdp.py b/torchacc/dist/fsdp.py
@@ -289,7 +289,7 @@ def full_optim_state_dict(model: torch.nn.Module,
                           xla_fsdp.XlaFullyShardedDataParallel) or isinstance(
                               model, FullyShardedDataParallel)
 
-        if isinstance(model, xla_fsdp.XlaFullyShardedDataParallel):
+        if isinstance(model, FullyShardedDataParallel):
             model = model.model
 
         shard_meta_data = model.get_shard_metadata()