def save()

in utils/helpers.py [0:0]

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

    def save(self, model_dir: str, suffix: str = ''):
        """
        :param model_dir: directory where the model should be stored
        :param suffix: option 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"
        th.save(self.state_dict(), fname)
        self.cuda()
        return self