# Buggy code
def __init__(
    self,
    base_lazy_tensor,
    left_interp_indices=None,
    left_interp_values=None,
    right_interp_indices=None,
    right_interp_values=None,
):
    base_lazy_tensor = lazify(base_lazy_tensor)

    if left_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        left_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        left_interp_indices.unsqueeze_(-1)
        left_interp_indices = left_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if left_interp_values is None:
        left_interp_values = torch.ones(
            left_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if right_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        right_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        right_interp_indices.unsqueeze_(-1)
        right_interp_indices = right_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if right_interp_values is None:
        right_interp_values = torch.ones(
            right_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if left_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:
        try:
            base_lazy_tensor = base_lazy_tensor._expand_batch(
                left_interp_indices.shape[:-2]
            )
        except RuntimeError:
            raise RuntimeError(
                "interp size ({}) is incompatible with base_lazy_tensor size ({}). ".format(
                    right_interp_indices.size(), base_lazy_tensor.size()
                )
            )

    super(InterpolatedLazyTensor, self).__init__(
        base_lazy_tensor,
        left_interp_indices,
        left_interp_values,
        right_interp_indices,
        right_interp_values,
    )
    self.base_lazy_tensor = base_lazy_tensor
    self.left_interp_indices = left_interp_indices
    self.left_interp_values = left_interp_values
    self.right_interp_indices = right_interp_indices
    self.right_interp_values = right_interp_values


# Traceback (RuntimeError)
(py37) vdhiman@dwarf:~/wrk/BayesCBF_ws/BayesCBF$ python tests/test_interpolated_lazy_tensor.py
Traceback (most recent call last):
File "tests/test_interpolated_lazy_tensor.py", line 30, in <module>
test_interpolated_lazy_tensor()
File "tests/test_interpolated_lazy_tensor.py", line 7, in test_interpolated_lazy_tensor
res = itplzt @ torch.eye(3)
File "/home/vdhiman/wrk/gpytorch/gpytorch/lazy/lazy_tensor.py", line 1731, in __matmul__
return self.matmul(other)
File "/home/vdhiman/wrk/gpytorch/gpytorch/lazy/interpolated_lazy_tensor.py", line 393, in matmul
right_interp_res = left_t_interp(self.right_interp_indices, self.right_interp_values, tensor, base_size)
File "/home/vdhiman/wrk/gpytorch/gpytorch/utils/interpolation.py", line 202, in left_t_interp
values = rhs.unsqueeze(-2) * interp_values.unsqueeze(-1)
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 0

# Fixed code
def __init__(
    self,
    base_lazy_tensor,
    left_interp_indices=None,
    left_interp_values=None,
    right_interp_indices=None,
    right_interp_values=None,
):
    base_lazy_tensor = lazify(base_lazy_tensor)

    if left_interp_indices is None:
        num_rows = base_lazy_tensor.size(-2)
        left_interp_indices = torch.arange(
            0, num_rows, dtype=torch.long, device=base_lazy_tensor.device
        )
        left_interp_indices.unsqueeze_(-1)
        left_interp_indices = left_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_rows, 1
        )

    if left_interp_values is None:
        left_interp_values = torch.ones(
            left_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if right_interp_indices is None:
        num_cols = base_lazy_tensor.size(-1)
        right_interp_indices = torch.arange(
            0, num_cols, dtype=torch.long, device=base_lazy_tensor.device
        )
        right_interp_indices.unsqueeze_(-1)
        right_interp_indices = right_interp_indices.expand(
            *base_lazy_tensor.batch_shape, num_cols, 1
        )

    if right_interp_values is None:
        right_interp_values = torch.ones(
            right_interp_indices.size(),
            dtype=base_lazy_tensor.dtype,
            device=base_lazy_tensor.device,
        )

    if left_interp_indices.shape[:-2] != base_lazy_tensor.batch_shape:
        try:
            base_lazy_tensor = base_lazy_tensor._expand_batch(
                left_interp_indices.shape[:-2]
            )
        except RuntimeError:
            raise RuntimeError(
                "interp size ({}) is incompatible with base_lazy_tensor size ({}). ".format(
                    right_interp_indices.size(), base_lazy_tensor.size()
                )
            )

    super(InterpolatedLazyTensor, self).__init__(
        base_lazy_tensor,
        left_interp_indices,
        left_interp_values,
        right_interp_indices,
        right_interp_values,
    )
    self.base_lazy_tensor = base_lazy_tensor
    self.left_interp_indices = left_interp_indices
    self.left_interp_values = left_interp_values
    self.right_interp_indices = right_interp_indices
    self.right_interp_values = right_interp_values