Saving and loading custom models

Hi, my custom model looks like this:

        mask_conv_model = models.resnet50(pretrained=True)
        for param in mask_conv_model.parameters():
            param.requires_grad = False
        num_ftrs = mask_conv_model.fc.in_features
        mask_conv_model.fc = nn.Sequential(nn.Linear(num_ftrs, masked_output_ftrs), nn.ReLU(), nn.Dropout())
        for name, param in mask_conv_model.named_parameters():
            if param.requires_grad == True:
                mask_params.append(param)
        mask_conv_model_optimizer = optim.Adam(mask_params, lr=masked_learning_rate)

Then later on, I save my model:

torch.save({
        'epoch': epoch,
        'target_seq_len': hype['target_seq_len'],
        'input_seq_len': hype['input_seq_len'],
        'masked_output_ftrs': hype['masked_output_ftrs'],
        'mask_conv': dict_of_models['MASKED CONV MODEL'].state_dict(),
    }, os.path.join(save_dir, "MASK_CONV.tar"))

Now, to load the model I have to write a lot of code:

        mask_conv_model = models.resnet50(pretrained=True)
        mask_checkpoint = torch.load(mask_dir)
        epoch = mask_checkpoint['epoch']
        masked_output_ftrs = mask_checkpoint['masked_output_ftrs']
        num_ftrs = mask_conv_model.fc.in_features
        mask_conv_model.fc = nn.Linear(num_ftrs, masked_output_ftrs)
        mask_conv_model.load_state_dict(mask_checkpoint['mask_conv'])

Is there a one-line method of loading models? Something like this:

mask_checkpoint = torch.load(mask_dir)
mask_conv_model = mask_checkpoint.load_state_dict(mask_checkpoint['mask_conv'])