# Buggy code
def forward(self, x1, x2, diag=False, **params):
    batch_shape = x1.shape[:-2]
    n_batch_dims = len(batch_shape)
    n1, d = x1.shape[-2:]
    n2 = x2.shape[-2]

    K = torch.zeros(
        *batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype
    )

    if not diag:
        # Scale the inputs by the lengthscale (for stability)
        x1_ = x1.div(self.lengthscale)
        x2_ = x2.div(self.lengthscale)

        # Form all possible rank-1 products for the gradient and Hessian blocks
        outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
        outer = outer / self.lengthscale.unsqueeze(-2)
        outer = torch.transpose(outer, -1, -2).contiguous()

        # 1) Kernel block
        diff = self.covar_dist(
            x1_, x2_, square_dist=True, dist_postprocess_func=postprocess_rbf, **params
        )
        K_11 = diff
        K[..., :n1, :n2] = K_11

        # 2) First gradient block
        outer1 = outer.view(*batch_shape, n1, n2 * d)
        K[..., :n1, n2:] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])

        # 3) Second gradient block
        outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
        outer2 = outer2.transpose(-1, -2)
        K[..., n1:, :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])

        # 4) Hessian block
        outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat(
            [*([1] * (n_batch_dims + 1)), d]
        )
        kp = KroneckerProductLazyTensor(
            torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1)
            / self.lengthscale.pow(2),
            torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(
                *batch_shape, 1, 1
            ),
        )
        chain_rule = kp.evaluate() - outer3
        K[..., n1:, n2:] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])

        # Symmetrize for stability
        if n1 == n2 and torch.eq(x1, x2).all():
            K = 0.5 * (K.transpose(-1, -2) + K)

        # Apply a perfect shuffle permutation to match the MutiTask ordering
        pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().reshape((n1 * (d + 1)))
        pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
        K = K[..., pi1, :][..., :, pi2]

        return K

    else:
        if not (n1 == n2 and torch.eq(x1, x2).all()):
            raise RuntimeError("diag=True only works when x1 == x2")

        kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
        grad_diag = torch.ones(
            *batch_shape, n2, d, device=x1.device, dtype=x1.dtype
        ) / self.lengthscale.pow_(2)
        grad_diag = grad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
        k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
        pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
        return k_diag[..., pi]


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-ecb6525006fe> in <module>
1 z = k(train_x, train_x, diag=True)
----> 2 z[0].backward()
3 print(k.raw_lengthscale.grad)

/mnt/xarfuse/uid-66331/be3771ae-seed-nspid4026531836-ns-4026531840/torch/tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
231                 create_graph=create_graph,
232                 inputs=inputs)
--> 233         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
234
235     def register_hook(self, hook):

/mnt/xarfuse/uid-66331/be3771ae-seed-nspid4026531836-ns-4026531840/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
144     Variable._execution_engine.run_backward(
145         tensors, grad_tensors_, retain_graph, create_graph, inputs,
--> 146         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
147
148

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1, 1]], which is output 0 of SoftplusBackward, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

# Fixed code
def forward(self, x1, x2, diag=False, **params):
    batch_shape = x1.shape[:-2]
    n_batch_dims = len(batch_shape)
    n1, d = x1.shape[-2:]
    n2 = x2.shape[-2]

    K = torch.zeros(
        *batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype
    )

    if not diag:
        # Scale the inputs by the lengthscale (for stability)
        x1_ = x1.div(self.lengthscale)
        x2_ = x2.div(self.lengthscale)

        # Form all possible rank-1 products for the gradient and Hessian blocks
        outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
        outer = outer / self.lengthscale.unsqueeze(-2)
        outer = torch.transpose(outer, -1, -2).contiguous()

        # 1) Kernel block
        diff = self.covar_dist(
            x1_, x2_, square_dist=True, dist_postprocess_func=postprocess_rbf, **params
        )
        K_11 = diff
        K[..., :n1, :n2] = K_11

        # 2) First gradient block
        outer1 = outer.view(*batch_shape, n1, n2 * d)
        K[..., :n1, n2:] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])

        # 3) Second gradient block
        outer2 = outer.transpose(-1, -3).reshape(*batch_shape, n2, n1 * d)
        outer2 = outer2.transpose(-1, -2)
        K[..., n1:, :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])

        # 4) Hessian block
        outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat(
            [*([1] * (n_batch_dims + 1)), d]
        )
        kp = KroneckerProductLazyTensor(
            torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1)
            / self.lengthscale.pow(2),
            torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(
                *batch_shape, 1, 1
            ),
        )
        chain_rule = kp.evaluate() - outer3
        K[..., n1:, n2:] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])

        # Symmetrize for stability
        if n1 == n2 and torch.eq(x1, x2).all():
            K = 0.5 * (K.transpose(-1, -2) + K)

        # Apply a perfect shuffle permutation to match the MutiTask ordering
        pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().reshape((n1 * (d + 1)))
        pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
        K = K[..., pi1, :][..., :, pi2]

        return K

    else:
        if not (n1 == n2 and torch.eq(x1, x2).all()):
            raise RuntimeError("diag=True only works when x1 == x2")

        kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
        grad_diag = torch.ones(
            *batch_shape, n2, d, device=x1.device, dtype=x1.dtype
        ) / self.lengthscale.pow(2)
        grad_diag = grad_diag.transpose(-1, -2).reshape(*batch_shape, n2 * d)
        k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
        pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().reshape((n2 * (d + 1)))
        return k_diag[..., pi]