def load()

in utils/helpers.py [0:0]

• 11 lines of code
• 2 McCabe index (conditional complexity)

    def load(self, model_dir: str, suffix: str = ''):
        """
        :param model_dir: directory where the model is be stored
        :param suffix: optional suffix to append to the network name
        """
        self.cpu()
        if suffix == "":
            fname = f"{model_dir}/{self.model_name}.pkl"
        else:
            fname = f"{model_dir}/{self.model_name}.{suffix}.pkl"
        states = th.load(fname)
        self.load_state_dict(states)
        self.cuda()
        print("Loaded:", fname)
        return self