Problem with @tf.function decoration: The tensor cannot be accessed from FuncGraph, because it was defined in FuncGraph which is out of scope

I have some code which runs in eager mode, which I have been debugging in order to get it to run in graph mode (decorated by tf.function). I am not sure about the error 'The tensor cannot be accessed from FuncGraph, because it was defined in FuncGraph which is out of scope.', which I have made MWE to reproduce here

import tensorflow as tf


class Outer(tf.keras.Model):
    def __init__(self, model):
        super().__init__()
        self.model = model
        self.mem = None
        self.loss_mem = None

    def loss(self, x):
        return 2*x-1

    @tf.function
    def gradient(self, x):
        self.mem = tf.TensorArray(tf.float32, size=10)
        self.loss_mem = tf.TensorArray(tf.float32, size=10)
        for i in tf.range(10):
            out_i = self.model.call(x)
            loss_i = self.loss(out_i)
            self.mem = self.mem.write(i, out_i)
            self.loss_mem = self.loss_mem.write(i, loss_i)

        ghat = [tf.zeros_like(i, dtype=tf.float32) for i in self.model.trainable_variables]
        for i in tf.reverse(tf.range(10), [0]):
            with tf.GradientTape() as g:
                g.watch(x)
                out_i = self.model.call(x)
            with tf.GradientTape() as g_loss:
                g_loss.watch(out_i)
                loss_i = self.loss(out_i)
            output_grad = g_loss.gradient(loss_i, out_i)
            ghat_update = g.gradient(out_i, self.model.trainable_variables, output_gradients=output_grad)
            for j in range(len(ghat)):
                ghat[j] = ghat[j] + ghat_update[j]
        return ghat

class Inner(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.out = tf.keras.layers.Dense(10)

    def call(self, x):
        return self.out(x[0])


x = [tf.zeros((32, 4, 16), tf.float32), tf.zeros((64, 10), tf.float32)]
mytest = Outer(Inner())
mytest.gradient(x)

I tried to reproduce with shorter examples but I cannot. The error is on this line

ghat[j] = ghat[j] + ghat_update[j]

I'm really at a loss for why this doesn't work, as I feel that ghat_update should definitely be in scope of the function.