# Buggy code
def loss_gradient(self, x, y, **kwargs):
    """
    Compute the gradient of the loss function w.r.t. `x`.

    :param x: Sample input with shape as expected by the model.
    :type x: `np.ndarray`
    :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or indices of shape
              (nb_samples,).
    :type y: `np.ndarray`
    :return: Array of gradients of the same shape as `x`.
    :rtype: `np.ndarray`
    """
    import mxnet as mx

    train_mode = self._learning_phase if hasattr(self, "_learning_phase") else False

    # Apply preprocessing
    x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=False)

    y_preprocessed = mx.nd.array([np.argmax(y_preprocessed, axis=1)]).T
    x_preprocessed = mx.nd.array(x_preprocessed.astype(ART_NUMPY_DTYPE), ctx=self._ctx)
    x_preprocessed.attach_grad()

    with mx.autograd.record(train_mode=train_mode):
        preds = self._model(x_preprocessed)
        loss = self._loss(preds, y_preprocessed)

    loss.backward()

    # Compute gradients
    grads = x_preprocessed.grad.asnumpy()
    grads = self._apply_preprocessing_gradient(x, grads)
    assert grads.shape == x.shape

    return grads


# Traceback (MXNetError)
MXNetError                                Traceback (most recent call last)
<timed exec> in <module>

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/attacks/attack.py in replacement_function(self, *args, **kwargs)
68                 if len(args) > 0:
69                     args = tuple(lst)
---> 70                 return fdict[func_name](self, *args, **kwargs)
71
72             replacement_function.__doc__ = fdict[func_name].__doc__

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/attacks/evasion/projected_gradient_descent.py in generate(self, x, y, **kwargs)
158                     self.eps_step,
159                     self._project,
--> 160                     self.num_random_init > 0 and i_max_iter == 0,
161                 )
162

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/attacks/evasion/fast_gradient.py in _compute(self, x, x_init, y, eps, eps_step, project, random_init)
318
319             # Get perturbation
--> 320             perturbation = self._compute_perturbation(batch, batch_labels)
321
322             # Apply perturbation and clip

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/attacks/evasion/fast_gradient.py in _compute_perturbation(self, batch, batch_labels)
273
274         # Get gradient wrt loss; invert it if attack is targeted
--> 275         grad = self.classifier.loss_gradient(batch, batch_labels) * (1 - 2 * int(self.targeted))
276
277         # Apply norm bound

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/classifiers/classifier.py in replacement_function(self, *args, **kwargs)
65                 if len(args) > 0:
66                     args = tuple(lst)
---> 67                 return fdict[func_name](self, *args, **kwargs)
68
69             replacement_function.__doc__ = fdict[func_name].__doc__

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/art/classifiers/mxnet.py in loss_gradient(self, x, y, **kwargs)
355         with mx.autograd.record(train_mode=train_mode):
356             preds = self._model(x_preprocessed)
--> 357             loss = self._loss(preds, y_preprocessed)
358
359         loss.backward()

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/gluon/block.py in __call__(self, *args)
691             hook(self, args)
692
--> 693         out = self.forward(*args)
694
695         for hook in self._forward_hooks.values():

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/gluon/block.py in forward(self, x, *args)
1156                     params = {k: v.data(ctx) for k, v in self._reg_params.items()}
1157
-> 1158                 return self.hybrid_forward(ndarray, x, *args, **params)
1159
1160         params = {i: j.var() for i, j in self._reg_params.items()}

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/gluon/loss.py in hybrid_forward(self, F, pred, label, sample_weight)
389             pred = log_softmax(pred, self._axis)
390         if self._sparse_label:
--> 391             loss = -pick(pred, label, axis=self._axis, keepdims=True)
392         else:
393             label = _reshape_like(F, label, pred)

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/ndarray/register.py in pick(data, index, axis, keepdims, mode, out, name, **kwargs)

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/_ctypes/ndarray.py in _imperative_invoke(handle, ndargs, keys, vals, out, is_np_op)
105         c_str_array(keys),
106         c_str_array([str(s) for s in vals]),
--> 107         ctypes.byref(out_stypes)))
108
109     create_ndarray_fn = _np_ndarray_cls if is_np_op else _ndarray_cls

~/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/base.py in check_call(ret)
253     """
254     if ret != 0:
--> 255         raise MXNetError(py_str(_LIB.MXGetLastError()))
256
257

MXNetError: [14:56:14] src/imperative/./imperative_utils.h:72: Check failed: inputs[i]->ctx().dev_mask() == ctx.dev_mask() (1 vs. 2) : Operator pick require all inputs live on the same context. But the first argument is on gpu(0) while the 2-th argument is on cpu(0)
Stack trace:
[bt] (0) /u/hesseltu/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/libmxnet.so(+0x6b8b5b) [0x2b4c5395bb5b]
[bt] (1) /u/hesseltu/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/libmxnet.so(mxnet::imperative::GetContext(nnvm::NodeAttrs const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, mxnet::Context const&amp;)+0x4fc) [0x2b4c56b3651c]
[bt] (2) /u/hesseltu/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/libmxnet.so(mxnet::Imperative::Invoke(mxnet::Context const&amp;, nnvm::NodeAttrs const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;)+0x1c0) [0x2b4c56b409a0]
[bt] (3) /u/hesseltu/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/libmxnet.so(+0x3759cef) [0x2b4c569fccef]
[bt] (4) /u/hesseltu/.pyenv/versions/3.7.6/envs/adversarialvideo/lib/python3.7/site-packages/mxnet/libmxnet.so(MXImperativeInvokeEx+0x62) [0x2b4c569fd2b2]
[bt] (5) /lib64/libffi.so.6(ffi_call_unix64+0x4c) [0x2b4c182a3dcc]
[bt] (6) /lib64/libffi.so.6(ffi_call+0x1f5) [0x2b4c182a36f5]
[bt] (7) /u/hesseltu/.pyenv/versions/3.7.6/lib/python3.7/lib-dynload/_ctypes.cpython-37m-x86_64-linux-gnu.so(_ctypes_callproc+0x283) [0x2b4c1808fa93]
[bt] (8) /u/hesseltu/.pyenv/versions/3.7.6/lib/python3.7/lib-dynload/_ctypes.cpython-37m-x86_64-linux-gnu.so(+0x8e3f) [0x2b4c18086e3f]

# Fixed code
def loss_gradient(self, x, y, **kwargs):
    """
    Compute the gradient of the loss function w.r.t. `x`.

    :param x: Sample input with shape as expected by the model.
    :type x: `np.ndarray`
    :param y: Target values (class labels) one-hot-encoded of shape (nb_samples, nb_classes) or indices of shape
              (nb_samples,).
    :type y: `np.ndarray`
    :return: Array of gradients of the same shape as `x`.
    :rtype: `np.ndarray`
    """
    import mxnet as mx

    train_mode = self._learning_phase if hasattr(self, "_learning_phase") else False

    # Apply preprocessing
    x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y, fit=False)

    y_preprocessed = mx.nd.array([np.argmax(y_preprocessed, axis=1)], ctx=self._ctx).T
    x_preprocessed = mx.nd.array(x_preprocessed.astype(ART_NUMPY_DTYPE), ctx=self._ctx)
    x_preprocessed.attach_grad()

    with mx.autograd.record(train_mode=train_mode):
        preds = self._model(x_preprocessed)
        loss = self._loss(preds, y_preprocessed)

    loss.backward()

    # Compute gradients
    grads = x_preprocessed.grad.asnumpy()
    grads = self._apply_preprocessing_gradient(x, grads)
    assert grads.shape == x.shape

    return grads