Error Loading a saved checkpoint outside model directory

callbacks
1

I have saved a model during training. However, I am facing issues when trying to load the model from a saved checkpoint using native pytorch outside the model directory. The model class name is ChangeNetSegment . On the python code torch.load('path/to/PyTLightning_checkpoint.pth') , I am getting the following error:

Traceback (most recent call last):
  File "pre_process_datasets/rename_changenet.py", line 6, in inspect_checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
  File "/home/user/miniconda3/envs/clip_new/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/user/miniconda3/envs/clip_new/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/user/miniconda3/envs/clip_new/lib/python3.8/site-packages/torch/serialization.py", line 1124, in find_class
    return super().find_class(mod_name, name)
ModuleNotFoundError: No module named 'module_name'

When I load the model keeping the relative directory structure same, I am able to successfully load the model. However, when I keep the module in ‘module_name’ directory and load the model from directory ‘B’, I am unable to load it.