PyTorch 2.6 has changed the default value of weights_only to True, breaking the loading of some nnunetv2 models

[joshuacwnewton]

Upstream context: pytorch/pytorch#52181, pytorch/pytorch#137602

PyTorch 2.6 isn't out on PyPI yet, but it just dropped on the http://download.pytorch.org/whl/cpu/torch/, and so we ran into the following class of error:

_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.optim.adam.Adam was not an allowed global by default. Please use `torch.serialization.add_safe_globals([Adam])` or the `torch.serialization.safe_globals([Adam])` context manager to allowlist this global if you trust this class/function.

Full list of errors (and the relevant globals) below:

 =========================== short test summary info ============================
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2s/t2s.nii.gz-/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2s/t2s_seg-deepseg.nii.gz-t2s_seg_deepseg.nii.gz-seg_sc_t2star-0.9] - _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 ivadomed.models.Unet was not an allowed global by default. Please use `torch.serialization.add_safe_globals([Unet])` or the `torch.serialization.safe_globals([Unet])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2_seg-manual.nii.gz-t2_seg_deepseg.nii.gz-seg_sc_contrast_agnostic-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.optim.adam.Adam was not an allowed global by default. Please use `torch.serialization.add_safe_globals([Adam])` or the `torch.serialization.safe_globals([Adam])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2_seg-deepseg_rootlets.nii.gz-t2_seg_deepseg.nii.gz-seg_spinal_rootlets_t2w-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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-None-t2_seg_deepseg.nii.gz-seg_sc_epi-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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-None-t2_seg_deepseg.nii.gz-seg_ms_lesion_mp2rage-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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_binary_seg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-None-t2_seg_deepseg.nii.gz-seg_ms_lesion-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 numpy._core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_softseg[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2_seg-manual.nii.gz-t2_seg_deepseg.nii.gz-seg_sc_contrast_agnostic-0] - _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.optim.adam.Adam was not an allowed global by default. Please use `torch.serialization.add_safe_globals([Adam])` or the `torch.serialization.safe_globals([Adam])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_softseg_error_with_fill_holes - _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.optim.adam.Adam was not an allowed global by default. Please use `torch.serialization.add_safe_globals([Adam])` or the `torch.serialization.safe_globals([Adam])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_multiclass[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2_fake_lesion.nii.gz-fnames_seg_manual0-t2_deepseg.nii.gz-suffixes0-seg_sc_lesion_t2w_sci-0.5] - _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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_multiclass[/__w/spinalcordtoolbox/spinalcordtoolbox/t2_ax.nii.gz-fnames_seg_manual1-t2_deepseg.nii.gz-suffixes1-seg_sc_ms_lesion_axial_t2w-0.5] - _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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.
FAILED testing/cli/test_sct_deepseg.py::test_segment_nifti_multiclass[/__w/spinalcordtoolbox/spinalcordtoolbox/data/sct_testing_data/t2/t2.nii.gz-fnames_seg_manual3-t2_deepseg.nii.gz-suffixes3-totalspineseg-0] - _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 numpy.core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([scalar])` or the `torch.serialization.safe_globals([scalar])` 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.

In our case, because the call to torch.load() occurs within nnunetv2, we can't set weights_only=False (even if we wanted to opt-in to the unsafe behavior).

That being said, I wanted to invite a larger discussion about how to make nnunetv2 models compatible with the new default of weights_only=True. I'm not intimately familiar with "safe globals" in PyTorch, or how nnunetv2 does (or does not) use them, so I would love to learn more about how nnunetv2 plans to handle this change.

Thanks much! ♥️

[wasserth]

I also have this issue. When you install pytorch (pip install torch) you can not use nnunet anymore. Because this will automatically install pytorch 2.6.
It seems that the easiest solution might be to simply set weights_only=False. A quick fix would be highly appreciated!

[wasserth]

I just tried it. If you change to

checkpoint = torch.load(join(model_training_output_dir, f'fold_{f}', checkpoint_name),
                                    map_location=torch.device('cpu'), weights_only=False)

in nnunetv2/inference/predict_from_raw_data.py then it works again.

[wasserth]

I created a PR with this fix.

[joshuacwnewton]

While setting weights_only=False does remove the error and preserve the older behavior, I imagine the torch devs upstream set weights_only=True as the new default for a good reason!

If nnunetv2 is known to universally set weights_only=False, then perhaps its users could be susceptible to some kind of arbitrary code execution attack due to a malicious model? (See: pytorch/pytorch#31875 and pytorch/pytorch#111806 for context.)

Of course there are pros and cons when it comes to maintenance expectations on the nnunetv2 devs, etc. but I imagine there should at least be some consideration taken? i.e. anything but a quick fix 😅. So, if you want a quick fix without troubling the nnunetv2 devs, then you should maybe set torch<2.6 for now, and wait to see what sort of conclusion they come to.

[snorthman]

The following fixes this issue:

import torch.serialization
import _codecs
torch.serialization.add_safe_globals([np.core.multiarray.scalar, np.dtype, np.dtypes.Float64DType, _codecs.encode])

just prior to running initialize_from_trained_model_folder.