InaccessibleTensorError with TensorArray

I have some code which runs in eager mode, I get InaccessibleTensorError when I run the code decorated with tf.function. Here is the error message

Traceback (most recent call last):
  File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\train_sde_jump_diffusion.py", line 55, in <module>
    losses = train(train_data, pastlen+predictlen, batch_size, num_batches, adjointgrad, optimizer)
  File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\train_sde_jump_diffusion.py", line 25, in train
    ghat = adjointgrad.gradient(cur_data)
  File "C:\Users\ronan\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow\python\util\traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\ronan\AppData\Local\Programs\Python\Python310\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1233, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.InaccessibleTensorError: in user code:

    File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\adjointgrad.py", line 153, in gradient  *
        model_return_loss = self.model.return_loss()
    File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\sde_jump_diffusion.py", line 123, in return_loss  *
        return self.loss_mem.stack(), self.x_mem.stack()[:self.predict_len]

    InaccessibleTensorError: <tf.Tensor 'while/TensorArrayV2Write/TensorListSetItem:0' shape=() dtype=variant> is out of scope and cannot be used here. Use return values, explicit Python locals or TensorFlow collections to access it.
    Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values for more information.
    
    <tf.Tensor 'while/TensorArrayV2Write/TensorListSetItem:0' shape=() dtype=variant> was defined here:
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\train_sde_jump_diffusion.py", line 55, in <module>
          losses = train(train_data, pastlen+predictlen, batch_size, num_batches, adjointgrad, optimizer)
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\train_sde_jump_diffusion.py", line 25, in train
          ghat = adjointgrad.gradient(cur_data)
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\adjointgrad.py", line 147, in gradient
          if self.need_check:
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\adjointgrad.py", line 151, in gradient
          for i in tf.range(n):
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\adjointgrad.py", line 152, in gradient
          self.model.call_and_sample(i)
        File "C:\Users\ronan\Documents\adjointgrad\adjointgrad\sde_jump_diffusion.py", line 114, in call_and_sample
          self.loss_mem = self.loss_mem.write(i, cur_loss)
    
    The tensor <tf.Tensor 'while/TensorArrayV2Write/TensorListSetItem:0' shape=() dtype=variant> cannot be accessed from FuncGraph(name=gradient, id=2370577038640), because it was defined in FuncGraph(name=while_body_1088, id=2370620445056), which is out of scope.

when I google these error messages, I get posts of people who aren't using TensorArray. But I'm already using TensorArray. The code is at github.com/ronan-keane/adjointgrad mostly in adjointgrad.py and sde_jump_diffusion.py, run train_sde_jump_diffusion to run. The code is basically like this

class AdjointGrad:
   def gradient(n):
      for i in tf.range(n):
         self.model.call_and_sample(i)
      self.model.return_loss()

class AdjointModel
   def call_and_sample(self, i):
      self.loss_mem = self.loss_mem.write(i, cur_loss)  # this is a tf.TensorArray