def compute_gradient(model_fn, x, y, targeted):
  """
  Computes the gradient of the loss with respect to the input tensor.
  :param model_fn: a callable that takes an input tensor and returns the model logits.
  :param x: input tensor
  :param y: Tensor with true labels. If targeted is true, then provide the target label.
  :param targeted:  bool. Is the attack targeted or untargeted? Untargeted, the default, will
                    try to make the label incorrect. Targeted will instead try to move in the
                    direction of being more like y.
  :return: A tensor containing the gradient of the loss with respect to the input tensor.
  """
  loss_fn = tf.nn.sparse_softmax_cross_entropy_with_logits
  with tf.GradientTape() as g:
    g.watch(x)
    # Compute loss
    loss = loss_fn(labels=y, logits=model_fn(x))
    if targeted:  # attack is targeted, minimize loss of target label rather than maximize loss of correct label
      loss = -loss

  # Define gradient of loss wrt input
  grad = g.gradient(loss, x)
  return grad