PyTorch 2.6.0 Release: Model Loading Fails with Unpickling Error

[SSWAMIN1SSS]

Error Observed:
With Latest version, Pytorch workload failed with error _pickle.UnpicklingError: Weights only load failed.

New version of torch(2.6.0) released on Jan 29th, has made change in the default behavior of the torch.load function. The default value for the weights_only parameter will be flipped to True

//./pytorchexample.py:8: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
 alexnet = torch.load("alexnet-pretrained.pt")

Steps to Reproduce

Follow the ReadMe steps as mentioned in the pytorch workload.
Issue is reproducible natively

• python3 pytorchexample.py

Actual Output:

Traceback (most recent call last):
  File "//./pytorchexample.py", line 8, in <module>
    alexnet = torch.load("alexnet-pretrained.pt")
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/intel/jenkins/workspace/local_ci_graphene_sgx_dcap/gramine/CI-Examples/pytorch/my_venv/lib/python3.12/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, �[1mdo those steps only if you trust the source of the checkpoint�[0m. 
	(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 torchvision.models.alexnet.AlexNet was not an allowed global by default. Please use `torch.serialization.add_safe_globals([AlexNet])` or the `torch.serialization.safe_globals([AlexNet])` 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.

Suggestions:
After adding weights_only=False pytorch example ran successfully.

diff --git a/pytorch/pytorchexample.py b/pytorch/pytorchexample.py

-alexnet = torch.load("alexnet-pretrained.pt")
+alexnet = torch.load("alexnet-pretrained.pt", weights_only=False)