Support PyTorch 2.4 load which requires explcit specification of weig…

…hts_only=False.

examples/air/main.py

@@ -200,7 +200,7 @@ def z_pres_prior_p(opt_step, time_step):
 
     if "load" in args:
         print("Loading parameters...")
-        air.load_state_dict(torch.load(args.load))
+        air.load_state_dict(torch.load(args.load, weights_only=False))
 
     # Viz sample from prior.
     if args.viz:

examples/cvae/cvae.py

@@ -184,6 +184,6 @@ def train(
             break
 
     # Save model weights
-    cvae_net.load_state_dict(torch.load(model_path))
+    cvae_net.load_state_dict(torch.load(model_path, weights_only=False))
     cvae_net.eval()
     return cvae_net

examples/dmm.py

@@ -465,7 +465,7 @@ def load_checkpoint():
             args.load_model
         ), "--load-model and/or --load-opt misspecified"
         logging.info("loading model from %s..." % args.load_model)
-        dmm.load_state_dict(torch.load(args.load_model))
+        dmm.load_state_dict(torch.load(args.load_model, weights_only=False))
         logging.info("loading optimizer states from %s..." % args.load_opt)
         adam.load(args.load_opt)
         logging.info("done loading model and optimizer states.")

pyro/contrib/examples/bart.py

@@ -11,6 +11,7 @@
 import subprocess
 import sys
 import urllib
+from functools import partial
 
 import torch
 
@@ -120,12 +121,12 @@ def load_bart_od():
         except urllib.error.HTTPError:
             logging.debug("cache miss, preprocessing from scratch")
     if os.path.exists(pkl_file):
-        return torch.load(pkl_file)
+        return torch.load(pkl_file, weights_only=False)
 
     filenames = multiprocessing.Pool(len(SOURCE_FILES)).map(
         _load_hourly_od, SOURCE_FILES
     )
-    datasets = list(map(torch.load, filenames))
+    datasets = list(map(partial(torch.load, weights_only=False), filenames))
 
     stations = sorted(set().union(*(d["stations"].keys() for d in datasets)))
     min_time = min(int(d["rows"][:, 0].min()) for d in datasets)

pyro/contrib/examples/nextstrain.py

@@ -41,4 +41,4 @@ def load_nextstrain_counts(map_location=None) -> dict:
     # Load tensors to the default location.
     if map_location is None:
         map_location = torch.tensor(0.0).device
-    return torch.load(filename, map_location=map_location)
+    return torch.load(filename, map_location=map_location, weights_only=False)

pyro/optim/optim.py

@@ -192,7 +192,9 @@ def load(self, filename: str, map_location=None) -> None:
         Load optimizer state from disk
         """
         with open(filename, "rb") as input_file:
-            state = torch.load(input_file, map_location=map_location)
+            state = torch.load(
+                input_file, map_location=map_location, weights_only=False
+            )
         self.set_state(state)
 
     def _get_optim(self, param: Union[Iterable[Tensor], Iterable[Dict[Any, Any]]]):

pyro/params/param_store.py

@@ -331,7 +331,7 @@ def load(self, filename: str, map_location: MAP_LOCATION = None) -> None:
         :type map_location: function, torch.device, string or a dict
         """
         with open(filename, "rb") as input_file:
-            state = torch.load(input_file, map_location)
+            state = torch.load(input_file, map_location, weights_only=False)
         self.set_state(state)
 
     @contextmanager

tests/contrib/cevae/test_cevae.py

@@ -64,7 +64,7 @@ def test_serialization(jit, feature_dim, outcome_dist):
             warnings.filterwarnings("ignore", category=UserWarning)
             torch.save(cevae, f)
         f.seek(0)
-        loaded_cevae = torch.load(f)
+        loaded_cevae = torch.load(f, weights_only=False)
 
     pyro.set_rng_seed(0)
     actual_ite = loaded_cevae.ite(x)

tests/contrib/easyguide/test_easyguide.py

@@ -89,7 +89,7 @@ def test_serialize():
         f = io.BytesIO()
         torch.save(guide, f)
         f.seek(0)
-        actual = torch.load(f)
+        actual = torch.load(f, weights_only=False)
 
     assert type(actual) == type(guide)
     assert dir(actual) == dir(guide)

tests/distributions/test_pickle.py

@@ -88,5 +88,5 @@ def test_pickle(Dist):
     # Note that pickling torch.Size() requires protocol >= 2
     torch.save(dist, buffer, pickle_protocol=pickle.HIGHEST_PROTOCOL)
     buffer.seek(0)
-    deserialized = torch.load(buffer)
+    deserialized = torch.load(buffer, weights_only=False)
     assert isinstance(deserialized, Dist)

tests/infer/mcmc/test_valid_models.py

@@ -420,7 +420,7 @@ def test_potential_fn_pickling(jit):
     buffer = io.BytesIO()
     torch.save(potential_fn, buffer)
     buffer.seek(0)
-    deser_potential_fn = torch.load(buffer)
+    deser_potential_fn = torch.load(buffer, weights_only=False)
     assert_close(deser_potential_fn(test_data), potential_fn(test_data))
 
 

tests/infer/test_autoguide.py

@@ -489,7 +489,7 @@ def test_serialization(auto_class, jit):
             f = io.BytesIO()
             torch.save(guide, f)
             f.seek(0)
-            guide_deser = torch.load(f)
+            guide_deser = torch.load(f, weights_only=False)
 
     # Check .call() result.
     pyro.set_rng_seed(0)

tests/nn/test_module.py

@@ -598,7 +598,7 @@ def test_mixin_factory():
         del module
         pyro.clear_param_store()
         f.seek(0)
-        module = torch.load(f)
+        module = torch.load(f, weights_only=False)
     assert type(module).__name__ == "PyroSequential"
     actual = module(data)
     assert_equal(actual, expected)
@@ -680,7 +680,7 @@ def test_torch_serialize_attributes(local_params):
             torch.save(module, f)
             pyro.clear_param_store()
             f.seek(0)
-            actual = torch.load(f)
+            actual = torch.load(f, weights_only=False)
 
         assert_equal(actual.x, module.x)
         actual_names = {name for name, _ in actual.named_parameters()}
@@ -704,7 +704,7 @@ def test_torch_serialize_decorators(local_params):
             torch.save(module, f)
             pyro.clear_param_store()
             f.seek(0)
-            actual = torch.load(f)
+            actual = torch.load(f, weights_only=False)
 
         assert_equal(actual.x, module.x)
         assert_equal(actual.y, module.y)

tests/poutine/test_poutines.py

@@ -1027,7 +1027,7 @@ def test_pickling(wrapper):
     # default protocol cannot serialize torch.Size objects (see https://github.com/pytorch/pytorch/issues/20823)
     torch.save(wrapped, buffer, pickle_protocol=pickle.HIGHEST_PROTOCOL)
     buffer.seek(0)
-    deserialized = torch.load(buffer)
+    deserialized = torch.load(buffer, weights_only=False)
     obs = torch.tensor(0.5)
     pyro.set_rng_seed(0)
     actual_trace = poutine.trace(deserialized).get_trace(obs)