InvalidArgumentError while using GRU layer in custom training loop
Author: dimitreOliveira
Pull request: no
Labels: name: wontfix

name: TensorFlow

**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

Comments:
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.