This is an internal helper function that can be used for either of these
    but not both at the same time:
    1. Record the output and gradient of output of an intermediate convolutional layer.
    2. Record the gradients of the image.

    Parameters
    ----------
    image : NDArray
        Image to visuaize. This is an NDArray with the preprocessed image.
    class_id : int
        Category ID this image belongs to. If not provided,
        network's prediction will be used.
    conv_layer_name: str
        Name of the convolutional layer whose output and output's gradients need to be acptured.
    image_grad: bool
        Whether to capture gradients of the image.

def _get_grad(net, image, class_id=None, conv_layer_name=None, image_grad=False):
    """This is an internal helper function that can be used for either of these
    but not both at the same time:
    1. Record the output and gradient of output of an intermediate convolutional layer.
    2. Record the gradients of the image.

    Parameters
    ----------
    image : NDArray
        Image to visuaize. This is an NDArray with the preprocessed image.
    class_id : int
        Category ID this image belongs to. If not provided,
        network's prediction will be used.
    conv_layer_name: str
        Name of the convolutional layer whose output and output's gradients need to be acptured.
    image_grad: bool
        Whether to capture gradients of the image."""

    if image_grad:
        image.attach_grad()
        Conv2D.capture_layer_name = None
        Activation.set_guided_backprop(True)
    else:
        # Tell convviz.Conv2D which layer's output and gradient needs to be recorded
        Conv2D.capture_layer_name = conv_layer_name
        Activation.set_guided_backprop(False)
    
    # Run the network
    with autograd.record(train_mode=False):
        out = net(image)
    
    # If user didn't provide a class id, we'll use the class that the network predicted
    if class_id == None:
        model_output = out.asnumpy()
        class_id = np.argmax(model_output)

    # Create a one-hot target with class_id and backprop with the created target
    one_hot_target = mx.nd.one_hot(mx.nd.array([class_id]), 1000)
    out.backward(one_hot_target, train_mode=False)

    if image_grad:
        return image.grad[0].asnumpy()
    else:
        # Return the recorded convolution output and gradient
        conv_out = Conv2D.conv_output
        return conv_out[0].asnumpy(), conv_out.grad[0].asnumpy()