Могу ли я использовать AutoGrad, если часть кода, который я хочу выделить, - это сеть PyTorch?

Question

Я пытаюсь использовать AutoGrad для вычисления производной определенного фрагмента кода.Часть этого кода состоит из нейронной сети, реализованной в PyTorch.Однако у меня есть некоторые проблемы с использованием AutoGrad для вычисления производной моего NN.

Я создал небольшой скрипт для воспроизведения проблемы:

import torch
import autograd.numpy as np
from autograd import grad

inputDimension = 10
hiddenLayerDimension = 10
outputDimension = 1

model = torch.nn.Sequential(
    torch.nn.Linear(inputDimension, hiddenLayerDimension),
    torch.nn.ReLU(),
    torch.nn.Linear(hiddenLayerDimension, outputDimension),
)

def functionToDifferentiate(input):
    # This line below represents the 'other' calculations. In reality it is more involved
    scaledInput = input * 3
    inputTensor = torch.from_numpy(scaledInput).type(torch.FloatTensor)
    return model(inputTensor)

randomInput = np.random.rand(inputDimension)
gradientFunctionOfModel = grad(functionToDifferentiate)

print(functionToDifferentiate(randomInput))
print(gradientFunctionOfModel(randomInput))

При запуске этого кода последняя строкавылетает со следующей трассировкой стека:

Traceback (most recent call last):
  File "replaced_for_privacy_reasons/stackOverFlowQuestion.py", line 25, in <module>
    print(gradientFunctionOfModel(randomInput))
  File "replaced_for_privacy_reasons/venv/lib/python3.6/site-packages/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
  File "replaced_for_privacy_reasons/venv/lib/python3.6/site-packages/autograd/differential_operators.py", line 24, in grad
    vjp, ans = _make_vjp(fun, x)
  File "replaced_for_privacy_reasons/venv/lib/python3.6/site-packages/autograd/core.py", line 10, in make_vjp
    end_value, end_node =  trace(start_node, fun, x)
  File "replaced_for_privacy_reasons/venv/lib/python3.6/site-packages/autograd/tracer.py", line 10, in trace
    end_box = fun(start_box)
tensor([0.0228], grad_fn=<AddBackward0>)
  File "replaced_for_privacy_reasons/venv/lib/python3.6/site-packages/autograd/wrap_util.py", line 15, in unary_f
    return fun(*subargs, **kwargs)
  File "replaced_for_privacy_reasons/stackOverFlowQuestion.py", line 18, in functionToDifferentiate
    inputTensor = torch.from_numpy(input).type(torch.FloatTensor)
TypeError: expected np.ndarray (got ArrayBox)

Я действительно знаю, что мои входы можно создавать непосредственно как тензоры PyTorch, а не как векторы Numpy, и вычислять градиенты относительно этих векторов.Однако это не очень хороший вариант для меня, поскольку эта нейронная сеть является лишь частью полной функции, для которой я хочу рассчитать градиент.

Любая помощь очень ценится