AttributeError: 'LogSoftmax' object has no attribute '_load_state_dict_pre_hooks'

1

I have torch version 1.0.1. When i am trying to load the checkpoint I am getting this error: AttributeError: ‘LogSoftmax’ object has no attribute ‘_load_state_dict_pre_hooks’.

The same load works fine with torch 0.4.0. Any help would be appreciated.

# TODO: Save the checkpoint
model.class_to_idx = train_image_datasets.class_to_idx
checkpoint_params = {'arch': 'vgg16',
                     'state_dict': model.state_dict(),
                     'class_to_idx': model.class_to_idx, 
                     'classifier':  model.classifier,
                     'input_s':input_feature,
                     'output_s':102,
                     'epochs':3}

torch.save(checkpoint_params, 'classifier_vgg16.pth')

and to load:

def load_checkpoint(filepath):
    #to avoid: RuntimeError: cuda runtime error (35) : 
    #CUDA driver version is insufficient for CUDA runtime version at torch/csrc/cuda/Module.cpp:51
    checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
    
    model = getattr(models, checkpoint['arch'])(pretrained=True)
    
    model.state_dict = checkpoint['state_dict']
    model.class_to_idx = checkpoint['class_to_idx']
    model.classifier = checkpoint['classifier']
    model.input_s = checkpoint['input_s']
    model.output_s = checkpoint['output_s']
    model.epochs = checkpoint['epochs']
    
    model.load_state_dict(checkpoint)
    return model

model = load_checkpoint(filepath='checkpoints/classifier_vgg16_new.pth')
model.to(device)