Checkpoint issue

[Xindi Zheng]

Hi,
You guys did an awesome work! I got some problem, when I run test.py. When I load 'trained_models/qm9_gaussian/checkpoints/model.ckpt' , I got error below

magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, 'v'.

[David Koes]

Where is the qm9_gaussian_ep directory coming from? I don't see it in the repository.

[Ian Dunn]

I realized the example command for sampling in the readme uses a checkpoint path that isn't in the repo. But, assuming you got around that issue, the cause of this exception could possibly be due to a mismatch of pickle versions? Or possibly a mismatch of PyTorch lightning versions?

Can you share the full stack trace you get when running the script. Can you confirm what version of python, torch, and PyTorch lightning you have installed in the environment using?

[Xindi Zheng]

Where is the qm9_gaussian_ep directory coming from? I don't see it in the repository.

Thanks for such a prompt response. Sorry, I provide the wrong path. I loaded 'trained_models/qm9_gaussian/checkpoints/model.ckpt'.

[Xindi Zheng]

Can you share the full stack trace you get when running the script. Can you confirm what version of python, torch, and PyTorch lightning you have installed in the environment using?

Hi Dr. Dunni3,
Thank you for answer.

That's my full stack trace:

  File "/home/xdzheng/sbdd/FlowMol/test.py", line 61, in <module>
    model = FlowMol.load_from_checkpoint(checkpoint_file)
  File "/home/xdzheng/miniconda3/envs/flowmol/lib/python3.10/site-packages/pytorch_lightning/core/module.py", line 1552, in load_from_checkpoint
    loaded = _load_from_checkpoint(
  File "/home/xdzheng/miniconda3/envs/flowmol/lib/python3.10/site-packages/pytorch_lightning/core/saving.py", line 61, in _load_from_checkpoint
    checkpoint = pl_load(checkpoint_path, map_location=map_location)
  File "/home/xdzheng/miniconda3/envs/flowmol/lib/python3.10/site-packages/lightning_fabric/utilities/cloud_io.py", line 55, in _load
    return torch.load(f, map_location=map_location)  # type: ignore[arg-type]
  File "/home/xdzheng/miniconda3/envs/flowmol/lib/python3.10/site-packages/torch/serialization.py", line 1028, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/home/xdzheng/miniconda3/envs/flowmol/lib/python3.10/site-packages/torch/serialization.py", line 1246, in _legacy_load
    magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, 'v'.

My conda environment:

python                    3.10.13
pytorch                   2.1.0
pytorch-lightning         2.1.0
pytorch-cuda              12.1

Best Regard~

[Xindi Zheng]

Dear all, Thanks for your answer. I fixed this issue by fully downloading the checkpoint. Really appreciate your excellent work.