InvalidArgumentError while using GRU layer in custom training loop
**System information**
- TensorFlow version `2.1.0`
- Python version: `3`
- GPU model and memory: `NVIDIA Tesla P100`
- CUDA Version: `10.1`
- Environment: This happens both on Kaggle and Colab

**Describe the current behavior**
I'm trying to train a Hugging face transformer model (roBERTa base) with a custom training loop, and got the error below:

```
InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  InstantiateOptions.input_devices must have the same length as the number of arguments: input_devices length = 23 number of arguments = 24
	 [[{{node while/body/_1/StatefulPartitionedCall}}]]
  (1) Invalid argument:  InstantiateOptions.input_devices must have the same length as the number of arguments: input_devices length = 23 number of arguments = 24
	 [[{{node while/body/_1/StatefulPartitionedCall}}]]
	 [[while/body/_1/Adam/Cast_6/ReadVariableOp/_30]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_step_35635]

Function call stack:
train_step -> train_step
```
The thing is I can run the same model using `model.fit()` API, and this error only happens when I use a LSTM or GRU layer on top of the transformer

**Describe the expected behavior**
Training should go normal
This is not enough information to help us. Can you post a minimal reproducible example, or at least how you construct the model. Also show where the error is triggered. (And just my two cents, but adding RNNs on top of transformer-based models seems redundant... But I guess you can try it out!)<|||||>Hi @BramVanroy , thanks for the tips, much appreciated, I was just trying different things to see how it performs, this is the model architecture:

```
module_config = RobertaConfig.from_pretrained(config['config_path'], output_hidden_states=False)

def model_fn(MAX_LEN):
    input_ids = layers.Input(shape=(MAX_LEN,), dtype=tf.int32, name='input_ids')
    attention_mask = layers.Input(shape=(MAX_LEN,), dtype=tf.int32, name='attention_mask')
    
    base_model = TFRobertaModel.from_pretrained(config['base_model_path'], config=module_config, name="base_model")
    last_hidden_state, _ = base_model({'input_ids': input_ids, 'attention_mask': attention_mask})
    
    x = layers.Bidirectional(layers.LSTM(128, return_sequences=True))(last_hidden_state)
    x = layers.Dropout(.1)(x)
    
    x_start = layers.TimeDistributed(layers.Dense(1))(x)
    x_start = layers.Flatten()(x_start)
    y_start = layers.Activation('softmax', name='y_start')(x_start)
    
    x_end = layers.TimeDistributed(layers.Dense(1))(x)
    x_end = layers.Flatten()(x_end)
    y_end = layers.Activation('softmax', name='y_end')(x_end)

    model = Model(inputs=[input_ids, attention_mask], outputs=[y_start, y_end])
    
    return model
```

And I was using it for a QA problem, so I also did:

```
model.compile(optimizer, loss={'y_start': losses.CategoricalCrossentropy(),
                                   'y_end': losses.CategoricalCrossentropy()})
```

The tricky part is that is jsut happens inside a custom training loop. Here are some of the code I have used.

```
# Step functions
    @tf.function
    def train_step(data_iter):
        def train_step_fn(x, y):
            with tf.GradientTape() as tape:
                probabilities = model(x, training=True)
                loss_start = loss_fn_start(y['y_start'], probabilities[0])
                loss_end = loss_fn_end(y['y_end'], probabilities[1])
                loss = tf.math.add(loss_start, loss_end)
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
            # update metrics
            train_acc_start.update_state(y['y_start'], probabilities)
            train_acc_end.update_state(y['y_end'], probabilities)
            train_loss.update_state(loss)
            train_loss_start.update_state(loss_start)
            train_loss_end.update_state(loss_end)
        for _ in tf.range(step_size):
            strategy.experimental_run_v2(train_step_fn, next(data_iter))

loss_fn_start = losses.categorical_crossentropy
loss_fn_end = losses.categorical_crossentropy

train_acc_start = metrics.CategoricalAccuracy()
train_acc_end = metrics.CategoricalAccuracy()
train_loss = metrics.Sum()
train_loss_start = metrics.Sum()
train_loss_end = metrics.Sum()
```

Let me know if you need any more information.<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
