Torch.load not able to find a class while importing a python script

data
1

I am trying to load a model for further processing, from a .pt save file after finishing the training. But while loading I getting the following error message:

Does this mean there has been some issue while saving the model? I saved the model with torch.save(unet, CONFIG['model_path2'])

While loading, I tried: unet_model = torch.load(f"{CONFIG['model_path2']}", map_location=torch.device(device=device)) and unet_model.load_state_dict(torch.load(f"{CONFIG['model_path2']}", map_location=torch.device(device=device))) but both are giving the same error message.

Is there any way I can check whether the saved .pt file has attributes for SinusoidalPosEmb or not?

Error:

Traceback (most recent call last):
  File "/Users/anshumansinha/Diff_structrepgen/gen_diff_1.py", line 67, in <module>
    unet_model.load_state_dict(torch.load(f"{CONFIG['model_path2']}", map_location=torch.device(device=device)))
  File "/Users/anshumansinha/miniconda3/lib/python3.10/site-packages/torch/serialization.py", line 809, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/Users/anshumansinha/miniconda3/lib/python3.10/site-packages/torch/serialization.py", line 1172, in _load
    result = unpickler.load()
  File "/Users/anshumansinha/miniconda3/lib/python3.10/site-packages/torch/serialization.py", line 1165, in find_class
    return super().find_class(mod_name, name)
AttributeError: Can't get attribute 'SinusoidalPosEmb' on <module '__main__' from '/Users/anshumansinha/Diff_structrepgen/gen_diff_1.py'>
2

It seems you have saved the entire model instead of the state_dict. Saving the model directly requires you to restore the same script structure when trying to load the model. Based on the error message I guess some scripts are missing or their definition changed.