ValueError in Interpolation with large samplesize/dimension

Hi! I was trying to use the library and ran into the error that reads "ValueError: left interp size (torch.Size([20000, 1, 1])) is incompatible with base lazy tensor size (torch.Size([20000, 20000])). Make sure the two have the same number of batch dimensions".

This only happens when I run on data with high-sample size/ dimension. I've modified the code in notebooks/bi_gp_ls.ipynb to replicate the error though the error occurs with different kernel settings. Is there an easy way to fix this, or are there extra steps when dealing with larger datasets? Thanks!

```
from tqdm.auto import tqdm
import torch
import gpytorch as gp
import altair as alt
import pandas as pd
import numpy as np

from gpytorch_lattice_kernel import RBFLattice as BilateralKernel

class BilateralGPModel(gp.models.ExactGP):
def __init__(self, train_x, train_y):
likelihood = gp.likelihoods.GaussianLikelihood()
super().__init__(train_x, train_y, likelihood)
# self.mean_module = gp.means.ConstantMean()
# self.covar_module = gp.kernels.ScaleKernel(BilateralKernel(ard_num_dims=train_x.size(-1)))
self.mean_module = gp.means.ZeroMean()
self.covar_module = BilateralKernel()

def forward(self, x):
mean_x = self.mean_module(x)
covar_x = self.covar_module(x)
return gp.distributions.MultivariateNormal(mean_x, covar_x)

def train(x, y, model, mll, optim):
model.train()

optim.zero_grad()

output = model(x)

loss = -mll(output, y)

loss.backward()

optim.step()

return { 'train/ll': -loss.detach().item() }

def test(x, y, model, lanc_iter=100, pre_size=0):
model.eval()

with torch.no_grad():
# gp.settings.max_preconditioner_size(pre_size), \
# gp.settings.max_root_decomposition_size(lanc_iter), \
# gp.settings.fast_pred_var():
preds = model(x)

pred_y = model.likelihood(model(x))
rmse = (pred_y.mean - y).pow(2).mean(0).sqrt()

return { 'test/rmse': rmse.item() }

def train_util(model, x, y, lr=0.1, epochs=100):
mll = gp.mlls.ExactMarginalLogLikelihood(model.likelihood, model)
optim = torch.optim.Adam(model.parameters(), lr=lr)

for _ in tqdm(range(epochs), leave=False):
train_dict = train(x, y, model, mll, optim)

return train_dict

n = 20000
d = 4
x = 2. * torch.rand(n, d) - 1.

with torch.no_grad():
covar_module = gp.kernels.ScaleKernel(gp.kernels.RBFKernel())
params = covar_module.state_dict()
params['raw_outputscale'] = torch.tensor(1.0).log()
params['base_kernel.raw_lengthscale'] = torch.Tensor([[1.5]]).log()
covar_module.load_state_dict(params)

covar = gp.distributions.MultivariateNormal(torch.zeros(n), covariance_matrix=covar_module(x))

rperm = torch.randperm(n)[:n//2]
train_x = x[rperm]
train_y = (covar.sample() + 0.1 * torch.randn(x.size(0)))[rperm]

for _ in tqdm(range(10)):
bigp = BilateralGPModel(train_x, train_y).float()

with gp.settings.max_root_decomposition_size(50):
train_dict = train_util(bigp, train_x, train_y)

for name, p in bigp.named_parameters():
results[name].append(p)
results['kind'].append('Bilateral GP')

for k, v in train_dict.items():
results[k].append(v)
```