[Tune/CI] Fix Tune-Pytorch-CIFAR notebook example (ray-project#26474)

Fixes the Tune-Pytorch-CIFAR notebook example as found in ray-project#26410 Signed-off-by: Antoni Baum <[email protected]> Signed-off-by: Stefan van der Kleij <[email protected]>
Stefan-1313 · Aug 18, 2022 · 3cb0848 · 3cb0848
1 parent 0e226e5
commit 3cb0848
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/doc/source/tune/examples/BUILD b/doc/source/tune/examples/BUILD
@@ -25,7 +25,6 @@ py_test_run_all_notebooks(
         "horovod_simple.ipynb", # REGRESSION
         "hyperopt_example.ipynb", # REGRESSION
         "sigopt_example.ipynb", # REGRESSION
-        "tune-pytorch-cifar.ipynb", # REGRESSION
         "tune-sklearn.ipynb", # REGRESSION
     ],
     data = ["//doc/source/tune/examples:tune_examples"],

diff --git a/doc/source/tune/examples/tune-pytorch-cifar.ipynb b/doc/source/tune/examples/tune-pytorch-cifar.ipynb
@@ -299,7 +299,7 @@
     "    device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
     "    best_trained_model.to(device)\n",
     "\n",
-    "    checkpoint_path = os.path.join(best_trial.checkpoint.dir_or_data, \"checkpoint\")\n",
+    "    checkpoint_path = os.path.join(best_trial.checkpoint.dir_or_data, \"checkpoint.pt\")\n",
     "\n",
     "    model_state, optimizer_state = torch.load(checkpoint_path)\n",
     "    best_trained_model.load_state_dict(model_state)\n",