Назначение переменных на месте с помощью pytorch

Question

Будучи (плохим) программистом с хорошими практиками, я пытаюсь написать общую "якобианскую" функцию в pytorch, которую я могу использовать для своего проекта.

#!/usr/bin/env python3
import numpy as np
import torch
import torch.autograd as ag

def full_jacobian(f, wrt):
    f_shape = list(f.size())
    wrt_shape = list(wrt.size())

    fj = ag.Variable(torch.zeros(f_shape + wrt_shape))


    f_range = nd_range(f_shape)
    wrt_range = nd_range(wrt_shape)

    for f_ind in f_range:
        fj[f_ind] = ag.grad(f[tuple(f_ind)], wrt, retain_graph=True, create_graph=True)[0]

Куда:

def nd_range(stop, dims = None):
    if dims == None:
        dims = len(stop)
    if not dims:
        yield ()
        return
    for outer in nd_range(stop, dims - 1):
        for inner in range(stop[dims - 1]):
            yield outer + (inner,)

К сожалению, когда я пытаюсь использовать это для вычисления якобиана, строка fj[f_ind] = ag.grad(f[tuple(f_ind)], wrt, retain_graph=True, create_graph=True)[0] перерывы с ошибкой:

RuntimeError: in-place operations can be only used on variables that don't share storage with any other variables, but detected that there are 2 objects sharing it.

Почему это ошибка? Какой правильный обходной путь?