Error with test the code

[shenweiyan]

I used the test code from the README as follows, and encountered a TypeError:

$ python3 predict.py -m models/kidney_1_1_CodonBert_model_20230726_320_model_param.pt -f data/example_data/protein_seq.fasta -o test/optimized.fasta
Traceback (most recent call last):
  File "predict.py", line 66, in <module>
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'), weights_only=False))
  File "/Bioinfo/Tools/Mamba/Mambaforge-24.3.0-0/envs/CodonBERT/lib/python3.7/site-packages/torch/serialization.py", line 594, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/Bioinfo/Tools/Mamba/Mambaforge-24.3.0-0/envs/CodonBERT/lib/python3.7/site-packages/torch/serialization.py", line 851, in _load
    unpickler = pickle_module.Unpickler(data_file, **pickle_load_args)
TypeError: 'weights_only' is an invalid keyword argument for Unpickler()

Due to server restrictions, many dependencies I used while installing CodonBERT are below the versions specified in the CodonBERT#installation instructions. Could this error be related to those dependencies?

Refer to the code below to use the conda environment for installation, but the same error still occurs.

conda env create -f codonbert_env.yaml -n codonbert_env
conda activate codonbert_env

$ python3 predict.py -m models/kidney_1_1_CodonBert_model_20230726_320_model_param.pt -f data/example_data/protein_seq.fasta -o test/optimized.fasta
Traceback (most recent call last):
  File "/home/bi.shenwy/software/CodonBERT/predict.py", line 66, in <module>
    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'), weights_only=False))
  File "/home/bi.shenwy/software/mambaforge-24.3.0-0/envs/CodonBERT/lib/python3.10/site-packages/torch/serialization.py", line 712, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/bi.shenwy/software/mambaforge-24.3.0-0/envs/CodonBERT/lib/python3.10/site-packages/torch/serialization.py", line 1047, in _load
    unpickler = UnpicklerWrapper(data_file, **pickle_load_args)
TypeError: 'weights_only' is an invalid keyword argument for Unpickler()

$ python3
Python 3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.__version__
'1.12.0'
>>>

[renzilin]

@ZhangDufei Please check it

[ZhangDufei]

Hi, @shenweiyan

The 'weights_only' option in torch.load() is only supported in PyTorch version >= 2.0

Try removing the weights_only=False parameter, which should resolve the issue:

model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

[shenweiyan]

Hi, @shenweiyan

The 'weights_only' option in torch.load() is only supported in PyTorch version >= 2.0

Try removing the weights_only=False parameter, which should resolve the issue:

model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

Thank you so much, it really solved the problem.