[Bug] Moving GP From CPU -> GPU after running it on CPU causes error #1619
Description
🐛 Bug
When I create a GP model, evaluate it on some data on the CPU, and then move the model to a GPU and try to evaluate it on more data it causes an error.
To reproduce
** Code snippet to reproduce **
import gpytorch
from gpytorch.models import ExactGP
from gpytorch.likelihoods import GaussianLikelihood
import torch
print(gpytorch.__version__)
class PlainGP(gpytorch.models.ExactGP):
def __init__(self, train_x, train_y, likelihood):
super().__init__(train_x, train_y, likelihood)
self.mean_module = gpytorch.means.ZeroMean()
self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
def forward(self, x):
mean = self.mean_module(x)
covar = self.covar_module(x)
return gpytorch.distributions.MultivariateNormal(mean, covar)
# Create training data
X_train = torch.linspace(-1, 1, 100).reshape(-1, 1)
y_train = torch.linspace(-1, 1, 100)
# Create model and likelihood
likelihood = GaussianLikelihood()
model = PlainGP(X_train, y_train, likelihood)
# Run model on data while on CPU
model.eval()
X_test = torch.linspace(-1, 1, 100).reshape(-1, 1)
preds = model(X_test)
# Move model to GPU and run model
model.cuda()
model(X_train.cuda())** Stack trace/error message **
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-37-c777d7a71d7b> in <module>
29 # Move model to GPU and run model
30 model.cuda()
---> 31 model(X_train.cuda())
/anaconda3/envs/base_hunter/lib/python3.8/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
317 # Make the prediction
318 with settings._use_eval_tolerance():
--> 319 predictive_mean, predictive_covar = self.prediction_strategy.exact_prediction(full_mean, full_covar)
320
321 # Reshape predictive mean to match the appropriate event shape
/anaconda3/envs/base_hunter/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_prediction(self, joint_mean, joint_covar)
260
261 return (
--> 262 self.exact_predictive_mean(test_mean, test_train_covar),
263 self.exact_predictive_covar(test_test_covar, test_train_covar),
264 )
/anaconda3/envs/base_hunter/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_predictive_mean(self, test_mean, test_train_covar)
278 # You **cannot* use addmv here, because test_train_covar may not actually be a non lazy tensor even for an exact
279 # GP, and using addmv requires you to delazify test_train_covar, which is obviously a huge no-no!
--> 280 res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
281 res = res + test_mean
282
RuntimeError: Tensor for argument #3 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for addmm)
It seems self.mean_cache is on CPU while test_train_covar is (correctly) on GPU
Expected Behavior
No error.
System information
Please complete the following information:
GPytorch Version: 1.4.1
PyTorch Version: 1.8.1
Ubuntu 18.04.5
Additional context
Add any other context about the problem here.