Dictionary error when training on OGB datasets

[AdarshMJ]

Hi!
I am trying to run the pruning for OGB datasets. I keep getting this error. Any idea how to resolve this? It runs for 1 epoch for 10 subgraphs and then I get this following error -

File "/home/relationalml/Desktop/Adarsh/GLT/Pruning/OGBN_proteins/unify/ogb/ogbn_proteins/main_imp.py", line 227, in <module>
    final_state_dict = main_get_mask(args, imp_num, resume_train_ckpt)
  File "/home/relationalml/Desktop/Adarsh/GLT/Pruning/OGBN_proteins/unify/ogb/ogbn_proteins/main_imp.py", line 178, in main_get_mask
    final_state_dict = pruning.save_all(dataset, 
  File "/home/relationalml/Desktop/Adarsh/GLT/Pruning/OGBN_proteins/unify/ogb/ogbn_proteins/pruning.py", line 73, in save_all
    'optimizer_state_dict': optimizer.state_dict()
  File "/home/relationalml/anaconda3/envs/gnnlottery/lib/python3.10/site-packages/torch/optim/optimizer.py", line 175, in state_dict
    packed_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v
  File "/home/relationalml/anaconda3/envs/gnnlottery/lib/python3.10/site-packages/torch/optim/optimizer.py", line 175, in <dictcomp>
    packed_state = {(param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v
KeyError: 140292191978176