Writing to TensorArray in tf.function - OperatorNotAllowedInGraphError

I am trying to get my code to run in graph mode (tf.function). I know it works in eager mode. In graph mode I get this error OperatorNotAllowedInGraphError: Using a symbolic `tf.Tensor` as a Python `bool` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature. The error is due to writing to the TensorArray. If anyone understands why this errors occurs, that would be helpful to know. This is related to my question from yesterday.

Here is a MWE for reproducing the error message

import tensorflow as tf


class AdjointGrad:
    def __init__(self, model):
        self.model = model

    @tf.function
    def gradient(self, n):
        for i in tf.range(n):
            self.model.call_and_sample(i)
        self.model.return_loss()

class AdjointModel:
    def __init__(self):
        self.loss_mem = tf.TensorArray(tf.float32, size=10)

    def call_and_sample(self, i):
        self.loss_mem = self.loss_mem.write(i, tf.ones((32,2)))

    def return_loss(self):
        return self.loss_mem.stack()

adjointgrad = AdjointGrad(AdjointModel())
adjointgrad.gradient(tf.cast(10, tf.int32))