Do not squeeze non-batch results in RootDecomposition backward
This caused issues when computing the gradients of a function computed
from the samples drawn from the posterior of a MTGP at a single point.
@@ -161,8 +161,6 @@ class RootDecomposition(Function):
right_factor = right_factor.contiguous()
res = lazy_tsr._quad_form_derivative(left_factor, right_factor)
- if not is_batch:
- res = [item.squeeze(0) for item in res]
return tuple(res)
else:
pass
