[Bug] LazyEvaluatedKernelTensor does not propagate requires_grad attribute

[trbedwards]

🐛 Bug

This is related to an issue I have filed in botorch in that gradients are not being computed when they should be.

I have a LazyEvaluatedKernelTensor which has the requires_grad=True attribute set. However, when calling .evaluate on this tensor, the resulting tensor has requires_grad=False and therefore is prevented from computing the gradient graph.

System information

• GPyTorch 1.6.0
• PyTorch 1.10.2
• Ubuntu 18.04.3 LTS 64-bit

[Balandat]

Can you provide a repro? The following works fine for me:

import torch
​from gpytorch.lazy import LazyEvaluatedKernelTensor
from gpytorch.kernels import RBFKernel

x1 = torch.rand(3, 2, requires_grad=True)  # this can be True or False since kernel parameters will make the LEKT require grad also
x2 = torch.rand(3, 2, requires_grad=True)
lekt = LazyEvaluatedKernelTensor(x1, x2, RBFKernel())

lekt.requires_grad  # True

out = lekt.evaluate()
out.requires_grad  # True

out.sum().backward()
x1.grad, x2.grad
(tensor([[ 2.1558, -0.5944],
         [ 1.4689,  1.7882],
         [ 0.2589,  1.0594]]),
 tensor([[-1.8151, -0.8231],
         [-0.3043, -1.0513],
         [-1.7642, -0.3787]]))