Saving Multi-GPU

[EoinKenny]

Thanks for this repo!

It seems like the code doesn't save the models correctly when using multiple GPUs right? Like if I train on multiple GPUs but load on a single CPU, it won't work I think.

You have this

    def save(self, path):
        states = {
            'G': self.G.state_dict(),
            'D': self.D.state_dict(),
            'optim_G': self.optim_G.state_dict(),
            'optim_D': self.optim_D.state_dict()
        }
        torch.save(states, path)

But I think you need self.G.module.state_dict() ?

[elvisyjlin]

Hello @EoinKenny , it could be the case when the model is trained on multiple GPUs. I trained the model on single GPU, so I didn't encounter any error. Please refine the code if you see it. Also, feel free to create a pull request!