Rename "model_state_dict" to "model"

bveeramani · Nov 22, 2022 · 5984f1e · 5984f1e
1 parent 5f54406
commit 5984f1e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py b/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py
@@ -83,7 +83,7 @@ def train_func(config):
         checkpoint_dict = session.get_checkpoint().to_dict()
 
         # Load in model
-        model_state = checkpoint_dict["model_state_dict"]
+        model_state = checkpoint_dict["model"]
         model.load_state_dict(model_state)
 
         # Load in optimizer
@@ -146,7 +146,7 @@ def train_func(config):
         checkpoint = Checkpoint.from_dict(
             {
                 "epoch": epoch,
-                "model_state_dict": model.state_dict(),
+                "model": model.state_dict(),
                 "optimizer_state_dict": optimizer.state_dict(),
             }
         )