From 9391f96f31e858a5a5a278383f24341636706c10 Mon Sep 17 00:00:00 2001 From: Davy Chen Date: Sat, 7 Sep 2024 20:20:34 +0800 Subject: [PATCH] # Fix 2 incorrect argument positioning 1. "weights_only" is augument of torch.load, not for model.load_state_dict Here is the doc: https://pytorch.org/docs/stable/generated/torch.load.html --- beginner_source/saving_loading_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/beginner_source/saving_loading_models.py b/beginner_source/saving_loading_models.py index 6c9b6b1fd77..13bca8ca3de 100644 --- a/beginner_source/saving_loading_models.py +++ b/beginner_source/saving_loading_models.py @@ -153,7 +153,7 @@ # .. code:: python # # model = TheModelClass(*args, **kwargs) -# model.load_state_dict(torch.load(PATH), weights_only=True) +# model.load_state_dict(torch.load(PATH, weights_only=True)) # model.eval() # # .. note:: @@ -407,7 +407,7 @@ # .. code:: python # # modelB = TheModelBClass(*args, **kwargs) -# modelB.load_state_dict(torch.load(PATH), strict=False, weights_only=True) +# modelB.load_state_dict(torch.load(PATH, weights_only=True), strict=False) # # Partially loading a model or loading a partial model are common # scenarios when transfer learning or training a new complex model.