How can I load model from checkpoints (with all params) for testing

I am trying to implement Unet for Carvana. I am taking help from a GitHub Repo.

I trained the model and save it in the checkpoint. However, when I am trying to test using the checkpoint model it is showing an error AttributeError: ‘dict’ object has no attribute ‘n_channels’

Getting error at the time of hparams

def main(hparams):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
net = Unet.load_from_checkpoint(hparams.checkpoint)

for fn in tqdm(os.listdir(hparams.img_dir)):
    fp = os.path.join(hparams.img_dir, fn)

    img =
    mask = predict(net, img, device=device)

    mask_img = mask_to_image(mask), fn))

if __name__ == '__main__':
parent_parser = ArgumentParser(add_help=False)
parent_parser.add_argument('--checkpoint', required=True)
parent_parser.add_argument('--img_dir', required=True)
parent_parser.add_argument('--out_dir', required=True)

parser = Unet.add_model_specific_args(parent_parser)
hparams = parser.parse_args()


Full code and

I am getting this error after loading the model, which one is related to n_channel

hparams are given from terminal python3 --checkpoint lightning_logs/version_0/checkpoints/_ckpt_epoch_1.ckpt --img_dir dataset/carvana/test --out_dir result/carvana

n_channels parameters are added from the using @staticmethod

Any kind of suggestion is applicable.

I solved the problem.

You have to use self.save_hyperparameters() in your Unet model to store the hyperparameters

I found the solution in this post here