import init_optim

pytorch · Dec 20, 2024 · 11e24f1 · 11e24f1
1 parent 91e9818
commit 11e24f1
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchtune/training/_distributed.py b/torchtune/training/_distributed.py
@@ -18,6 +18,7 @@
 from torch.distributed._tensor import distribute_tensor, DTensor
 from torch.distributed._tensor.placement_types import DTensorSpec, TensorMeta
 from torch.distributed.checkpoint.state_dict import (
+    _init_optim_state,
     get_model_state_dict,
     get_optimizer_state_dict,
     set_model_state_dict,