# Buggy code
def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
    # pass in the values we saved automatically
    if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
        model_args = {}

        # add some back compatibility, the actual one shall be last
        for hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS + (
            cls.CHECKPOINT_HYPER_PARAMS_KEY,
        ):
            if hparam_key in checkpoint:
                model_args.update(checkpoint[hparam_key])

        model_args = _convert_loaded_hparams(
            model_args, checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_TYPE)
        )

        args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)
        cls_spec = inspect.getfullargspec(cls.__init__)
        kwargs_identifier = cls_spec.varkw
        cls_init_args_name = inspect.signature(cls).parameters.keys()

        if args_name == "kwargs":
            # in case the class cannot take any extra argument filter only the possible
            if not kwargs_identifier:
                model_args = {
                    k: v for k, v in model_args.items() if k in cls_init_args_name
                }
            cls_kwargs.update(**model_args)
        elif args_name:
            if args_name in cls_init_args_name:
                cls_kwargs.update({args_name: model_args})
        else:
            cls_args = (model_args,) + cls_args

    # load the state_dict on the model automatically
    model = cls(*cls_args, **cls_kwargs)
    model.load_state_dict(checkpoint["state_dict"])

    # give model a chance to load something
    model.on_load_checkpoint(checkpoint)

    return model


# Traceback (TypeError)
TypeError                                 Traceback (most recent call last)

<ipython-input-5-50449ee4f6cc> in <module>()
1 test_loader = DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32)
----> 2 trainer.test(test_dataloaders=test_loader)

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in test(self, model, test_dataloaders, ckpt_path)
1168             if ckpt_path == 'best':
1169                 ckpt_path = self.checkpoint_callback.best_model_path
-> 1170             model = self.get_model().load_from_checkpoint(ckpt_path)
1171
1172         self.testing = True

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/core/saving.py in load_from_checkpoint(cls, checkpoint_path, map_location, hparams_file, tags_csv, *args, **kwargs)
167         checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
168
--> 169         model = cls._load_model_state(checkpoint, *args, **kwargs)
170         return model
171

/usr/local/lib/python3.6/dist-packages/pytorch_lightning/core/saving.py in _load_model_state(cls, checkpoint, *cls_args, **cls_kwargs)
201
202         # load the state_dict on the model automatically
--> 203         model = cls(*cls_args, **cls_kwargs)
204         model.load_state_dict(checkpoint['state_dict'])
205

TypeError: __init__() takes 1 positional argument but 2 were given

# Fixed code
def _load_model_state(cls, checkpoint: Dict[str, Any], *cls_args, **cls_kwargs):
    cls_spec = inspect.getfullargspec(cls.__init__)
    cls_init_args_name = inspect.signature(cls).parameters.keys()
    # pass in the values we saved automatically
    if cls.CHECKPOINT_HYPER_PARAMS_KEY in checkpoint:
        model_args = {}

        # add some back compatibility, the actual one shall be last
        for hparam_key in CHECKPOINT_PAST_HPARAMS_KEYS + (
            cls.CHECKPOINT_HYPER_PARAMS_KEY,
        ):
            if hparam_key in checkpoint:
                model_args.update(checkpoint[hparam_key])

        model_args = _convert_loaded_hparams(
            model_args, checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_TYPE)
        )

        args_name = checkpoint.get(cls.CHECKPOINT_HYPER_PARAMS_NAME)

        if args_name == "kwargs":
            # in case the class cannot take any extra argument filter only the possible
            cls_kwargs.update(**model_args)
        elif args_name:
            if args_name in cls_init_args_name:
                cls_kwargs.update({args_name: model_args})
        else:
            cls_args = (model_args,) + cls_args

    if not cls_spec.varkw:
        # filter kwargs according to class init unless it allows any argument via kwargs
        cls_kwargs = {k: v for k, v in cls_kwargs.items() if k in cls_init_args_name}

    # prevent passing positional arguments if class does not accept any
    if len(cls_spec.args) <= 1 and not cls_spec.kwonlyargs:
        cls_args, cls_kwargs = [], {}
    model = cls(*cls_args, **cls_kwargs)
    # load the state_dict on the model automatically
    model.load_state_dict(checkpoint["state_dict"])

    # give model a chance to load something
    model.on_load_checkpoint(checkpoint)

    return model