f3dabcf2c16d8ba0c702b18c9fccb633fa8b90d5,pycox/torch_fitter.py,FitNet,save_model_weights,#FitNet#,99

Before Change

            path: The filepath of the model.
            **kwargs: Arguments passed to torch.save method.
        """
        torch.save(self.net.state_dict(), path, **kwargs)

    def load_model_weights(self, path, **kwargs):
        """Load model weights.

After Change

            path: The filepath of the model.
            **kwargs: Arguments passed to torch.save method.
        """
        if self.cuda is not False:
            // If we don"t do this, torch will move g to default gpu.
            // Remove this when bug is fixed...
            self.net.cpu()
            torch.save(self.net.state_dict(), path, **kwargs)
            to_cuda(self.net, self.cuda)
        else:
            torch.save(self.net.state_dict(), path, **kwargs)

    def load_model_weights(self, path, **kwargs):
        """Load model weights.

        Parameters:

Instances