_pickle.UnpicklingError: Weights only load failed when freeze a model

[wangshuai-simulation]

when i execute: dp --pt freeze -c model.pt -o model_frozen.pth
here is the error:
To get the best performance, it is recommended to adjust the number of threads by setting the environment variables OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. See https://deepmd.rtfd.io/parallelism/ for more information.
[2025-04-17 23:33:42,660] DEEPMD INFO DeePMD version: 3.0.2
Traceback (most recent call last):
File "/home/shuaiwang/anaconda3/envs/deepmdkit/bin/dp", line 8, in
sys.exit(main())
^^^^^^
File "/home/shuaiwang/anaconda3/envs/deepmdkit/lib/python3.12/site-packages/deepmd/main.py", line 927, in main
deepmd_main(args)
File "/home/shuaiwang/anaconda3/envs/deepmdkit/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 355, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/home/shuaiwang/anaconda3/envs/deepmdkit/lib/python3.12/site-packages/deepmd/pt/entrypoints/main.py", line 547, in main
freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
File "/home/shuaiwang/anaconda3/envs/deepmdkit/lib/python3.12/site-packages/deepmd/pt/entrypoints/main.py", line 368, in freeze
model = inference.Tester(model, head=head).model
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/shuaiwang/anaconda3/envs/deepmdkit/lib/python3.12/site-packages/deepmd/pt/infer/inference.py", line 37, in init
state_dict = torch.load(model_ckpt, map_location=DEVICE, weights_only=True)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/shuaiwang/anaconda3/envs/deepmdkit/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 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.

how to avoid the error? Appreciate if anyone can help me !

[wangshuai-simulation]

Here the model.pt is from zenodo of the article of 'DPA-2: A Large Atomic Model As a Multi-task learner 'https://zenodo.org/records/13342300
The version of model.pt is v4 and in the downloaded file of Code_model_script/Model/DPA2_18head_1m

[njzjz]

The model in this article is only compatible with the code in this article.