TPU error: Input 2 to node `CTCLoss/ctc-loss/ctc_state_trans/ScatterNd_1` must be a compile-time constant

[Alex Flint]

Thanks for the all the work that has gone into this project. I tried to run the code from the fine-tuning colab notebook on a TPU but I get the following error when running model.fit:

  (0) INVALID_ARGUMENT: {{function_node __inference_train_function_84983}} Input 2 to node `CTCLoss/ctc-loss/ctc_state_trans/ScatterNd_1` with op ScatterNd must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. This error means that a shape or dimension argument could not be evaluated at compile time, usually because the value of the argument depends on a parameter to the computation, on a variable, or on a stateful operation such as a random number generator.

As per the colab notebook I am loading a pretrained wav2vec2 from tfhub::

  load_locally = tf.saved_model.LoadOptions(experimental_io_device='/job:localhost')  # required for TPU
  pretrained_layer = tfhub.KerasLayer("https://tfhub.dev/vasudevgupta7/wav2vec2/1", trainable=True, load_options=load_locally)

(I added "load_locally" to make it tfhub.KerasLayer work on the TPU)

I guess the pretrained model uses the scatter_nd op in a way that is not compatible with the TPU, right? Any idea what I can do about this?

I did see your end-to-end training script, by the way, which I see has TPU support, but I was hoping to load a pretrained model.