# Buggy code
def interpolate(self, x_grid, x_target, interp_points=range(-2, 2)):
    # Do some boundary checking
    grid_mins = x_grid.min(0)[0]
    grid_maxs = x_grid.max(0)[0]
    x_target_min = x_target.min(0)[0]
    x_target_max = x_target.min(0)[0]
    lt_min_mask = (x_target_min - grid_mins).lt(-1e-7)
    gt_max_mask = (x_target_max - grid_maxs).gt(1e-7)
    if lt_min_mask.sum().item():
        first_out_of_range = lt_min_mask.nonzero().squeeze(1)[0].item()
        raise RuntimeError(
            (
                "Received data that was out of bounds for the specified grid. "
                "Grid bounds were ({0:.3f}, {0:.3f}), but min = {0:.3f}, "
                "max = {0:.3f}"
            ).format(
                grid_mins[first_out_of_range].item(),
                grid_maxs[first_out_of_range].item(),
                x_target_min[first_out_of_range].item(),
                x_target_max[first_out_of_range].item(),
            )
        )
    if gt_max_mask.sum().item():
        first_out_of_range = gt_max_mask.nonzero().squeeze(1)[0].item()
        raise RuntimeError(
            (
                "Received data that was out of bounds for the specified grid. "
                "Grid bounds were ({0:.3f}, {0:.3f}), but min = {0:.3f}, "
                "max = {0:.3f}"
            ).format(
                grid_mins[first_out_of_range].item(),
                grid_maxs[first_out_of_range].item(),
                x_target_min[first_out_of_range].item(),
                x_target_max[first_out_of_range].item(),
            )
        )

    # Now do interpolation
    interp_points = torch.tensor(
        interp_points, dtype=x_grid.dtype, device=x_grid.device
    )
    interp_points_flip = interp_points.flip(0)

    num_grid_points = x_grid.size(0)
    num_target_points = x_target.size(0)
    num_dim = x_target.size(-1)
    num_coefficients = len(interp_points)

    interp_values = torch.ones(
        num_target_points,
        num_coefficients**num_dim,
        dtype=x_grid.dtype,
        device=x_grid.device,
    )
    interp_indices = torch.zeros(
        num_target_points,
        num_coefficients**num_dim,
        dtype=torch.long,
        device=x_grid.device,
    )

    for i in range(num_dim):
        grid_delta = x_grid[1, i] - x_grid[0, i]
        lower_grid_pt_idxs = torch.floor(
            (x_target[:, i] - x_grid[0, i]) / grid_delta
        ).squeeze()
        lower_pt_rel_dists = (
            x_target[:, i] - x_grid[0, i]
        ) / grid_delta - lower_grid_pt_idxs
        lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
        lower_grid_pt_idxs.detach_()

        scaled_dist = lower_pt_rel_dists.unsqueeze(-1) + interp_points_flip.unsqueeze(
            -2
        )
        dim_interp_values = self._cubic_interpolation_kernel(scaled_dist)

        # Find points who's closest lower grid point is the first grid point
        # This corresponds to a boundary condition that we must fix manually.
        left_boundary_pts = torch.nonzero(lower_grid_pt_idxs < 1)
        num_left = len(left_boundary_pts)

        if num_left > 0:
            left_boundary_pts.squeeze_(1)
            x_grid_first = (
                x_grid[:num_coefficients, i]
                .unsqueeze(1)
                .t()
                .expand(num_left, num_coefficients)
            )

            grid_targets = (
                x_target.select(1, i)[left_boundary_pts]
                .unsqueeze(1)
                .expand(num_left, num_coefficients)
            )
            dists = torch.abs(x_grid_first - grid_targets)
            closest_from_first = torch.min(dists, 1)[1]

            for j in range(num_left):
                dim_interp_values[left_boundary_pts[j], :] = 0
                dim_interp_values[left_boundary_pts[j], closest_from_first[j]] = 1
                lower_grid_pt_idxs[left_boundary_pts[j]] = 0

        right_boundary_pts = torch.nonzero(
            lower_grid_pt_idxs > num_grid_points - num_coefficients
        )
        num_right = len(right_boundary_pts)

        if num_right > 0:
            right_boundary_pts.squeeze_(1)
            x_grid_last = (
                x_grid[-num_coefficients:, i]
                .unsqueeze(1)
                .t()
                .expand(num_right, num_coefficients)
            )

            grid_targets = x_target.select(1, i)[right_boundary_pts].unsqueeze(1)
            grid_targets = grid_targets.expand(num_right, num_coefficients)
            dists = torch.abs(x_grid_last - grid_targets)
            closest_from_last = torch.min(dists, 1)[1]

            for j in range(num_right):
                dim_interp_values[right_boundary_pts[j], :] = 0
                dim_interp_values[right_boundary_pts[j], closest_from_last[j]] = 1
                lower_grid_pt_idxs[right_boundary_pts[j]] = (
                    num_grid_points - num_coefficients
                )

        offset = (interp_points - interp_points.min()).long().unsqueeze(-2)
        dim_interp_indices = lower_grid_pt_idxs.long().unsqueeze(-1) + offset

        n_inner_repeat = num_coefficients**i
        n_outer_repeat = num_coefficients ** (num_dim - i - 1)
        index_coeff = num_grid_points ** (num_dim - i - 1)
        dim_interp_indices = dim_interp_indices.unsqueeze(-1).repeat(
            1, n_inner_repeat, n_outer_repeat
        )
        dim_interp_values = dim_interp_values.unsqueeze(-1).repeat(
            1, n_inner_repeat, n_outer_repeat
        )
        interp_indices = interp_indices.add(
            dim_interp_indices.view(num_target_points, -1).mul(index_coeff)
        )
        interp_values = interp_values.mul(dim_interp_values.view(num_target_points, -1))

    return interp_indices, interp_values


# Traceback (IndexError)
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-14-dd2c9e04445d> in <module>()
1 # Does not work
2 test_x = torch.FloatTensor([-0.1]).unsqueeze(1)
----> 3 model(test_x)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
142                 n_train=n_train,
143                 likelihood=self.likelihood,
--> 144                 precomputed_cache=self.mean_cache,
145             )
146             predictive_covar, covar_cache = exact_predictive_covar(

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/functions/__init__.py in exact_predictive_mean(full_covar, full_mean, train_labels, n_train, likelihood, precomputed_cache)
89
90         full_covar = NonLazyVariable(full_covar)
---> 91     return full_covar.exact_predictive_mean(full_mean, train_labels, n_train, likelihood, precomputed_cache)
92
93

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in exact_predictive_mean(self, full_mean, train_labels, n_train, likelihood, precomputed_cache)
124         else:
125             return super(LazyEvaluatedKernelVariable, self).exact_predictive_mean(
--> 126                 full_mean, train_labels, n_train, likelihood, precomputed_cache
127             )
128

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in exact_predictive_mean(self, full_mean, train_labels, n_train, likelihood, precomputed_cache)
427         else:
428             test_train_covar = self[n_train:, :n_train]
--> 429         res = test_train_covar.matmul(precomputed_cache)
430         if res.ndimension() == 3:
431             res = res.squeeze(-1)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_variable.py in matmul(self, tensor)
632             raise RuntimeError
633
--> 634         func = Matmul(self.representation_tree())
635         return func(tensor, *self.representation())
636

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in representation_tree(self)
112
113     def representation_tree(self):
--> 114         return LazyVariableRepresentationTree(self.evaluate_kernel())
115
116     def evaluate(self):

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/lazy/lazy_evaluated_kernel_variable.py in evaluate_kernel(self)
96                 x2 = self.x2
97
---> 98             self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
99             if self.squeeze_row:
100                 self._cached_kernel_eval.squeeze_(-2)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
160
161     def __call__(self, *inputs, **kwargs):
--> 162         outputs = self.forward(*inputs, **kwargs)
163         if torch.is_tensor(outputs) or isinstance(outputs, RandomVariable) or isinstance(outputs, LazyVariable):
164             return outputs

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/multitask_kernel.py in forward(self, x1, x2)
61     def forward(self, x1, x2):
62         covar_i = self.task_covar_module.covar_matrix
---> 63         covar_x = self.data_covar_module.forward(x1, x2)
64         if covar_x.size(0) == 1:
65             covar_x = covar_x[0]

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/grid_interpolation_kernel.py in forward(self, x1, x2, **kwargs)
56             base_lazy_var = base_lazy_var.repeat(x1.size(0), 1, 1)
57
---> 58         left_interp_indices, left_interp_values = self._compute_grid(x1)
59         if torch.equal(x1.data, x2.data):
60             right_interp_indices = left_interp_indices

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/kernels/grid_interpolation_kernel.py in _compute_grid(self, inputs)
39         batch_size, n_data, n_dimensions = inputs.size()
40         inputs = inputs.view(batch_size * n_data, n_dimensions)
---> 41         interp_indices, interp_values = Interpolation().interpolate(Variable(self.grid), inputs)
42         interp_indices = interp_indices.view(batch_size, n_data, -1)
43         interp_values = interp_values.view(batch_size, n_data, -1)

~/anaconda/envs/py3/lib/python3.6/site-packages/gpytorch/utils/interpolation.py in interpolate(self, x_grid, x_target, interp_points)
114                     dim_interp_values[left_boundary_pts[j], :] = 0
115                     dim_interp_values[left_boundary_pts[j], closest_from_first[j]] = 1
--> 116                     lower_grid_pt_idxs[left_boundary_pts[j]] = 0
117
118             right_boundary_pts = torch.nonzero(lower_grid_pt_idxs > num_grid_points - num_coefficients)

IndexError: too many indices for tensor of dimension 0

# Fixed code
def interpolate(self, x_grid, x_target, interp_points=range(-2, 2)):
    # Do some boundary checking
    grid_mins = x_grid.min(0)[0]
    grid_maxs = x_grid.max(0)[0]
    x_target_min = x_target.min(0)[0]
    x_target_max = x_target.min(0)[0]
    lt_min_mask = (x_target_min - grid_mins).lt(-1e-7)
    gt_max_mask = (x_target_max - grid_maxs).gt(1e-7)
    if lt_min_mask.sum().item():
        first_out_of_range = lt_min_mask.nonzero().squeeze(1)[0].item()
        raise RuntimeError(
            (
                "Received data that was out of bounds for the specified grid. "
                "Grid bounds were ({0:.3f}, {0:.3f}), but min = {0:.3f}, "
                "max = {0:.3f}"
            ).format(
                grid_mins[first_out_of_range].item(),
                grid_maxs[first_out_of_range].item(),
                x_target_min[first_out_of_range].item(),
                x_target_max[first_out_of_range].item(),
            )
        )
    if gt_max_mask.sum().item():
        first_out_of_range = gt_max_mask.nonzero().squeeze(1)[0].item()
        raise RuntimeError(
            (
                "Received data that was out of bounds for the specified grid. "
                "Grid bounds were ({0:.3f}, {0:.3f}), but min = {0:.3f}, "
                "max = {0:.3f}"
            ).format(
                grid_mins[first_out_of_range].item(),
                grid_maxs[first_out_of_range].item(),
                x_target_min[first_out_of_range].item(),
                x_target_max[first_out_of_range].item(),
            )
        )

    # Now do interpolation
    interp_points = torch.tensor(
        interp_points, dtype=x_grid.dtype, device=x_grid.device
    )
    interp_points_flip = interp_points.flip(0)

    num_grid_points = x_grid.size(0)
    num_target_points = x_target.size(0)
    num_dim = x_target.size(-1)
    num_coefficients = len(interp_points)

    interp_values = torch.ones(
        num_target_points,
        num_coefficients**num_dim,
        dtype=x_grid.dtype,
        device=x_grid.device,
    )
    interp_indices = torch.zeros(
        num_target_points,
        num_coefficients**num_dim,
        dtype=torch.long,
        device=x_grid.device,
    )

    for i in range(num_dim):
        grid_delta = x_grid[1, i] - x_grid[0, i]
        lower_grid_pt_idxs = torch.floor(
            (x_target[:, i] - x_grid[0, i]) / grid_delta
        ).squeeze()
        lower_pt_rel_dists = (
            x_target[:, i] - x_grid[0, i]
        ) / grid_delta - lower_grid_pt_idxs
        lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
        lower_grid_pt_idxs.detach_()

        if len(lower_grid_pt_idxs.shape) == 0:
            lower_grid_pt_idxs = lower_grid_pt_idxs.unsqueeze(0)

        scaled_dist = lower_pt_rel_dists.unsqueeze(-1) + interp_points_flip.unsqueeze(
            -2
        )
        dim_interp_values = self._cubic_interpolation_kernel(scaled_dist)

        # Find points who's closest lower grid point is the first grid point
        # This corresponds to a boundary condition that we must fix manually.
        left_boundary_pts = torch.nonzero(lower_grid_pt_idxs < 1)
        num_left = len(left_boundary_pts)

        if num_left > 0:
            left_boundary_pts.squeeze_(1)
            x_grid_first = (
                x_grid[:num_coefficients, i]
                .unsqueeze(1)
                .t()
                .expand(num_left, num_coefficients)
            )

            grid_targets = (
                x_target.select(1, i)[left_boundary_pts]
                .unsqueeze(1)
                .expand(num_left, num_coefficients)
            )
            dists = torch.abs(x_grid_first - grid_targets)
            closest_from_first = torch.min(dists, 1)[1]

            for j in range(num_left):
                dim_interp_values[left_boundary_pts[j], :] = 0
                dim_interp_values[left_boundary_pts[j], closest_from_first[j]] = 1
                lower_grid_pt_idxs[left_boundary_pts[j]] = 0

        right_boundary_pts = torch.nonzero(
            lower_grid_pt_idxs > num_grid_points - num_coefficients
        )
        num_right = len(right_boundary_pts)

        if num_right > 0:
            right_boundary_pts.squeeze_(1)
            x_grid_last = (
                x_grid[-num_coefficients:, i]
                .unsqueeze(1)
                .t()
                .expand(num_right, num_coefficients)
            )

            grid_targets = x_target.select(1, i)[right_boundary_pts].unsqueeze(1)
            grid_targets = grid_targets.expand(num_right, num_coefficients)
            dists = torch.abs(x_grid_last - grid_targets)
            closest_from_last = torch.min(dists, 1)[1]

            for j in range(num_right):
                dim_interp_values[right_boundary_pts[j], :] = 0
                dim_interp_values[right_boundary_pts[j], closest_from_last[j]] = 1
                lower_grid_pt_idxs[right_boundary_pts[j]] = (
                    num_grid_points - num_coefficients
                )

        offset = (interp_points - interp_points.min()).long().unsqueeze(-2)
        dim_interp_indices = lower_grid_pt_idxs.long().unsqueeze(-1) + offset

        n_inner_repeat = num_coefficients**i
        n_outer_repeat = num_coefficients ** (num_dim - i - 1)
        index_coeff = num_grid_points ** (num_dim - i - 1)
        dim_interp_indices = dim_interp_indices.unsqueeze(-1).repeat(
            1, n_inner_repeat, n_outer_repeat
        )
        dim_interp_values = dim_interp_values.unsqueeze(-1).repeat(
            1, n_inner_repeat, n_outer_repeat
        )
        interp_indices = interp_indices.add(
            dim_interp_indices.view(num_target_points, -1).mul(index_coeff)
        )
        interp_values = interp_values.mul(dim_interp_values.view(num_target_points, -1))

    return interp_indices, interp_values