Whenever i tried to save the entire model the warning shows, it saves without any errors.
checkpoint = {'model': ResNet9(), 'state_dict': model.state_dict()}
torch.save(checkpoint, 'model.pth')

But when i try loading it , it gives me this error.
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-1-29ed9644f643> in <module>
     10     return model
     11 
---> 12 model = load_checkpoint('face.pth')

<ipython-input-1-29ed9644f643> in load_checkpoint(filepath)
      1 import torch
      2 def load_checkpoint(filepath):
----> 3     checkpoint = torch.load(filepath)
      4     model = checkpoint['model']
      5     model.load_state_dict(checkpoint['state_dict'])

~/anaconda3/lib/python3.7/site-packages/torch/serialization.py in load(f, map_location, pickle_module, **pickle_load_args)
    578                     return torch.jit.load(f)
    579                 return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
--> 580         return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
    581 
    582 

~/anaconda3/lib/python3.7/site-packages/torch/serialization.py in _legacy_load(f, map_location, pickle_module, **pickle_load_args)
    758     unpickler = pickle_module.Unpickler(f, **pickle_load_args)
    759     unpickler.persistent_load = persistent_load
--> 760     result = unpickler.load()
    761 
    762     deserialized_storage_keys = pickle_module.load(f, **pickle_load_args)

AttributeError: Can't get attribute 'ResNet9' on <module '__main__'>

my torch version is 1.5.1 and im using colab