DeepSSMUtils -- torch.load fails

[siboles]

Description

PyTorch has changed the default value for the weights_only argument in torch.load. This causes an error using the dataloaders.

Error Trace

Traceback (most recent call last):
  File "/home/scotts/shapeworks/ShapeWorks-v6.5.1-linux/Examples/Python/RunUseCase.py", line 97, in <module>
    module.Run_Pipeline(args)
  File "/home/scotts/shapeworks/ShapeWorks-v6.5.1-linux/Examples/Python/deep_ssm.py", line 459, in Run_Pipeline
    predicted_val_world_particles = DeepSSMUtils.testDeepSSM(
  File "/home/scotts/miniforge3/envs/shapeworks/lib/python3.9/site-packages/DeepSSMUtils/__init__.py", line 51, in testDeepSSM
    predicted_particle_files = eval.test(config_file, loader)
  File "/home/scotts/miniforge3/envs/shapeworks/lib/python3.9/site-packages/DeepSSMUtils/eval.py", line 38, in test
    test_loader = torch.load(loader_dir + loader)
  File "/home/scotts/miniforge3/envs/shapeworks/lib/python3.9/site-packages/torch/serialization.py", line 1470, in load
    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint.
	(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
	(2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
	WeightsUnpickler error: Unsupported global: GLOBAL torch.utils.data.dataloader.DataLoader was not an allowed global by default. Please use `torch.serialization.add_safe_globals([DataLoader])` or the `torch.serialization.safe_globals([DataLoader])` context manager to allowlist this global if you trust this class/function.

Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

A Fix

I chose option (1) in from the error trace recommendation. Adding weights_only=False to all occurrences of torch.load in DeepSSMUtils.eval, DeepSSMUtils.model, and DeepSSMUtils.trainer resolved the issue. As PyTorch warns, this will expose a vulnerability to execute arbitrary code though.

I'd be happy to submit this simple pull request, but another solution may be more ideal.

[akenmorris]

What version of PyTorch are you using? The conda environment setup by ShapeWorks should have 1.11.0 or 1.13.1

[siboles]

Yes, sorry I should have mentioned that. My system CUDA is 12.6, so I replaced the CPU version installed by the shell script without thinking to pin the version

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126

which fetches '2.6.0+cu126'.

That said, if you do plan to upgrade to a newer torch later, I haven't encountered any issues so far after the minor fix. Feel free to close this issue.