Question about saving StyleTransfer model files

[zanxiR]

Hi! I know that your model file only saves parameters in the end. Have you ever tried to save a complete network model file? I used the following code in trainer.py to save and kept reporting errors.

    def save(self, iterations):
        gen_name = os.path.join(self.model_dir, 'gen_%08d.pt' % (iterations + 1))
        dis_name = os.path.join(self.model_dir, 'dis_%08d.pt' % (iterations + 1))
        opt_name = os.path.join(self.model_dir, 'optimizer.pt')
        torch.save({'gen': self.model.gen.state_dict()}, gen_name)
        torch.save({'dis': self.model.dis.state_dict()}, dis_name)
        torch.save({'gen': self.gen_opt.state_dict(),
                    'dis': self.dis_opt.state_dict()}, opt_name)
        torch.save(self.model.gen, '\modelgen.pt')
        torch.save(self.model.dis, '\modeldis.pt')

error： AttributeError: Can't pickle local object 'boolean_dispatch.<locals>.fn' I searched for this error on the Internet, but I haven't found a solution yet.

If you have a pth file with complete parameters, is it convenient to share it? Thank you!