Running PyKeOps in Google Colab encounter TypeError with LazyTensor

[NicPy4]

I tried implementing the SpectralMixtureKernel to KeOps because I handle big data and experienced pretty good extrapolation on smaller data sets without KeOps. I'm using Google Colab with the T4 GPU (15GB). While testing my code I encounter a

TypeError: object of class LazyTensor cannot be made into a Tensor

I thought that maybe I did something wrong with the implementation of the SpectralMixtureKernel, so I copied the example from https://docs.gpytorch.ai/en/latest/examples/02_Scalable_Exact_GPs/KeOps_GP_Regression.html. However, it throws the same error again.

TypeError                                 Traceback (most recent call last)
[<ipython-input-14-7e02bc603789>](https://localhost:8080/#) in <cell line: 13>()
     18     output = model(train_x)
     19     # Calc loss and backprop gradients
---> 20     loss = -mll(output, train_y)
     21     loss.backward()
     22     print('Iter %d/%d - Loss: %.3f   lengthscale: %.3f   noise: %.3f' % (

12 frames
[/usr/local/lib/python3.10/dist-packages/gpytorch/module.py](https://localhost:8080/#) in __call__(self, *inputs, **kwargs)
     29 
     30     def __call__(self, *inputs, **kwargs) -> Union[Tensor, Distribution, LinearOperator]:
---> 31         outputs = self.forward(*inputs, **kwargs)
     32         if isinstance(outputs, list):
     33             return [_validate_module_outputs(output) for output in outputs]

[/usr/local/lib/python3.10/dist-packages/gpytorch/mlls/exact_marginal_log_likelihood.py](https://localhost:8080/#) in forward(self, function_dist, target, *params)
     62         # Get the log prob of the marginal distribution
     63         output = self.likelihood(function_dist, *params)
---> 64         res = output.log_prob(target)
     65         res = self._add_other_terms(res, params)
     66 

[/usr/local/lib/python3.10/dist-packages/gpytorch/distributions/multivariate_normal.py](https://localhost:8080/#) in log_prob(self, value)
    191         # Get log determininant and first part of quadratic form
    192         covar = covar.evaluate_kernel()
--> 193         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    194 
    195         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/_linear_operator.py](https://localhost:8080/#) in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
   1746             args = [inv_quad_rhs] + list(args)
   1747 
-> 1748         preconditioner, precond_lt, logdet_p = self._preconditioner()
   1749         if precond_lt is None:
   1750             from ..operators.identity_linear_operator import IdentityLinearOperator

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/added_diag_linear_operator.py](https://localhost:8080/#) in _preconditioner(self)
    124         if self._q_cache is None:
    125             max_iter = settings.max_preconditioner_size.value()
--> 126             self._piv_chol_self = self._linear_op.pivoted_cholesky(rank=max_iter)
    127             if torch.any(torch.isnan(self._piv_chol_self)).item():
    128                 warnings.warn(

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/_linear_operator.py](https://localhost:8080/#) in pivoted_cholesky(self, rank, error_tol, return_pivots)
   1963         """
   1964         func = PivotedCholesky.apply
-> 1965         res, pivots = func(self.representation_tree(), rank, error_tol, *self.representation())
   1966 
   1967         if return_pivots:

[/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py](https://localhost:8080/#) in apply(cls, *args, **kwargs)
    537             # See NOTE: [functorch vjp and autograd interaction]
    538             args = _functorch.utils.unwrap_dead_wrappers(args)
--> 539             return super().apply(*args, **kwargs)  # type: ignore[misc]
    540 
    541         if cls.setup_context == _SingleLevelFunction.setup_context:

[/usr/local/lib/python3.10/dist-packages/linear_operator/functions/_pivoted_cholesky.py](https://localhost:8080/#) in forward(ctx, representation_tree, max_iter, error_tol, *matrix_args)
     22         # Need to get diagonals. This is easy if it's a LinearOperator, since
     23         # LinearOperator.diagonal() operates in batch mode.
---> 24         matrix_diag = matrix._approx_diagonal()
     25         # NOTE: we will be performing inpace operations on matrix_diag
     26         # Calling _approx_diag() may return a tensor that shares the same storage as the LinearOperator

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/constant_mul_linear_operator.py](https://localhost:8080/#) in _approx_diagonal(self)
     72 
     73     def _approx_diagonal(self: Float[LinearOperator, "*batch N N"]) -> Float[torch.Tensor, "*batch N"]:
---> 74         res = self.base_linear_op._approx_diagonal()
     75         return res * self._constant.unsqueeze(-1)
     76 

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/_linear_operator.py](https://localhost:8080/#) in _approx_diagonal(self)
    490         :return: the (batch of) diagonals (... x N)
    491         """
--> 492         return self._diagonal()
    493 
    494     @cached(name="cholesky")

[/usr/local/lib/python3.10/dist-packages/linear_operator/utils/memoize.py](https://localhost:8080/#) in g(self, *args, **kwargs)
     57         kwargs_pkl = pickle.dumps(kwargs)
     58         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
---> 59             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
     60         return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
     61 

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/kernel_linear_operator.py](https://localhost:8080/#) in _diagonal(self)
    231         x2 = self.x2.unsqueeze(0).transpose(0, -2)
    232         tensor_params = {name: val.unsqueeze(0) for name, val in self.tensor_params.items()}
--> 233         diag_mat = to_dense(self.covar_func(x1, x2, **tensor_params, **self.nontensor_params))
    234         assert diag_mat.shape[-2:] == torch.Size(self.num_outputs_per_input)
    235 

[/usr/local/lib/python3.10/dist-packages/linear_operator/operators/_linear_operator.py](https://localhost:8080/#) in to_dense(obj)
   2985         return obj.to_dense()
   2986     else:
-> 2987         raise TypeError("object of class {} cannot be made into a Tensor".format(obj.__class__.__name__))
   2988 
   2989 

TypeError: object of class LazyTensor cannot be made into a Tensor

So now I was wondering if there is any issue with Google Colab or why it throws this error.