TFBert: Out of memory error when acting on a strided slice of input.
Author: codeninja
Pull request: no
Labels: name: wontfix

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): TF Bert & TFBertForSequenceEncoding

Language I am using the model on (English, Chinese ...): EN

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [X] my own task or dataset: (give details below)

## Code To reproduce

```
batch_size = 10
window = 50


# seq_len = 512 # not good to override unless the sequence truncatting is disabled.

from transformers import TFBertModel, TFBertForSequenceClassification, BertTokenizer
# configuration = BertConfig()

def build_model(num_labels = 10, max_seq_len=seq_len, use_logits=True, batch_size = batch_size, label_size = None):
    print(f"Number of labels: {num_labels}")

    optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-06, clipnorm=1.0)
    loss = tf.keras.losses.CategoricalCrossentropy(from_logits=use_logits)
    macc = tf.keras.metrics.CategoricalAccuracy('accuracy')
    print(f"Optimizer used: {optimizer}")
    print(f"Loss used: {loss}")
    print(f"Acc used: {macc}")
    
    num_labels = len(unique_classes) # 9 unique classes
    bert_config = BertConfig.from_pretrained("bert-base-cased", 
                                         num_labels=len(unique_classes), 
                                         output_hidden_states=False, 
                                         output_attentions=True)
    
    bert_model = TFBertForSequenceClassification.from_pretrained("bert-base-cased", config=bert_config)


    
    # [ ENCODINGS, ATTN_MASK]
    agent_input = Input(shape=[2, seq_len], batch_size=batch_size, name='agent_input', dtype='int32')
    print("Agent Input:",agent_input)

    print('Stacking Inputs')
    agent_encodings = agent_input[...,0,:]
    agent_attn = agent_input[...,1,:]
    print('Agent Encodings:', agent_encodings)
    print('Agent Attn: ', agent_attn)
    
    print("Building Bert Model")
    agent_outputs = bert_model([agent_encodings, agent_attn])
    agent_predictions = agent_outputs[0]
    agent_attn = agent_outputs[1]
    
    print('agent_outputs', agent_outputs)
    model = Model(inputs=agent_input, outputs=agent_predictions)
    model.compile(optimizer=optimizer, loss=loss, metrics=[macc])


    return model

#_________________________


print('shape for train data', agent_encodings_train.shape)
batch_size = 100
bert_model = build_model(len(unique_classes), seq_len, True, batch_size)
bert_model.summary()

#_________________________

bert_model.evaluate(agent_encodings_train[:100], agent_class_train[:100], batch_size=batch_size)
```
![image](https://user-images.githubusercontent.com/14914/80419128-2aa6b480-889e-11ea-8a91-a7025a07cedc.png)

![image](https://user-images.githubusercontent.com/14914/80419229-5cb81680-889e-11ea-89b1-96f051fa58d6.png)

agent_encodings_train contains a (batch_size, 2, sequence_len) shape of tensors containing seq_len encoded bert tokens and seq_len bert attention tokens.

I am slicing them in the model into each token part and passing them to bert. You will notice in the summary. there are 2 strided slices being passed to BERT. The model compiles, and evaluates, but when calling model.fit I constantly run into memory problems as seen below.

```
bert_model.summary()
history = bert_model.fit(agent_encodings_train[:100], agent_class_train[:100],
                epochs=1,
                batch_size=batch_size,
                class_weights=class_weights,
                shuffle=False)
```

```
---------------------------------------------------------------------------
ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-115-712b9a078eba> in <module>
     10 #                 validation_data=(agent_encodings_test, agent_class_test),
     11                 class_weights=class_weights,
---> 12                 shuffle=False)

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\keras\engine\training.py in _method_wrapper(self, *args, **kwargs)
     63   def _method_wrapper(self, *args, **kwargs):
     64     if not self._in_multi_worker_mode():  # pylint: disable=protected-access
---> 65       return method(self, *args, **kwargs)
     66 
     67     # Running inside `run_distribute_coordinator` already.

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\keras\engine\training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_batch_size, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)
    781                 batch_size=batch_size):
    782               callbacks.on_train_batch_begin(step)
--> 783               tmp_logs = train_function(iterator)
    784               # Catch OutOfRangeError for Datasets of unknown size.
    785               # This blocks until the batch has finished executing.

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\def_function.py in __call__(self, *args, **kwds)
    577         xla_context.Exit()
    578     else:
--> 579       result = self._call(*args, **kwds)
    580 
    581     if tracing_count == self._get_tracing_count():

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\def_function.py in _call(self, *args, **kwds)
    641         # Lifting succeeded, so variables are initialized and we can run the
    642         # stateless function.
--> 643         return self._stateless_fn(*args, **kwds)
    644     else:
    645       canon_args, canon_kwds = \

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\function.py in __call__(self, *args, **kwargs)
   2418     with self._lock:
   2419       graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
-> 2420     return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
   2421 
   2422   @property

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\function.py in _filtered_call(self, args, kwargs)
   1663          if isinstance(t, (ops.Tensor,
   1664                            resource_variable_ops.BaseResourceVariable))),
-> 1665         self.captured_inputs)
   1666 
   1667   def _call_flat(self, args, captured_inputs, cancellation_manager=None):

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\function.py in _call_flat(self, args, captured_inputs, cancellation_manager)
   1744       # No tape is watching; skip to running the function.
   1745       return self._build_call_outputs(self._inference_function.call(
-> 1746           ctx, args, cancellation_manager=cancellation_manager))
   1747     forward_backward = self._select_forward_and_backward_functions(
   1748         args,

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\function.py in call(self, ctx, args, cancellation_manager)
    596               inputs=args,
    597               attrs=attrs,
--> 598               ctx=ctx)
    599         else:
    600           outputs = execute.execute_with_cancellation(

~\Anaconda3\envs\cat2_nightly\lib\site-packages\tensorflow\python\eager\execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     58     ctx.ensure_initialized()
     59     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 60                                         inputs, attrs, num_outputs)
     61   except core._NotOkStatusException as e:
     62     if name is not None:

ResourceExhaustedError: 2 root error(s) found.
  (0) Resource exhausted:  OOM when allocating tensor with shape[100,86,3072] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[node model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/activation/truediv (defined at C:\Users\codeninja\Anaconda3\envs\cat2_nightly\lib\site-packages\transformers\modeling_tf_bert.py:63) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

	 [[clip_by_norm_2/truediv/_30]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

  (1) Resource exhausted:  OOM when allocating tensor with shape[100,86,3072] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[node model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/activation/truediv (defined at C:\Users\codeninja\Anaconda3\envs\cat2_nightly\lib\site-packages\transformers\modeling_tf_bert.py:63) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_253267]

Errors may have originated from an input operation.
Input Source operations connected to node model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/activation/truediv:
 model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/dense/BiasAdd (defined at C:\Users\codeninja\Anaconda3\envs\cat2_nightly\lib\site-packages\transformers\modeling_tf_bert.py:319)

Input Source operations connected to node model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/activation/truediv:
 model_11/tf_bert_for_sequence_classification_31/bert/encoder/layer_._0/intermediate/dense/BiasAdd (defined at C:\Users\codeninja\Anaconda3\envs\cat2_nightly\lib\site-packages\transformers\modeling_tf_bert.py:319)

Function call stack:
train_function -> train_function
```


<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->

## Environment info
<!-- You can run the command `transformers-cli env
` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

[env.txt](https://github.com/huggingface/transformers/files/4542124/env.txt)

- `transformers` version:
- Platform:
- Python version:
- PyTorch version (GPU?):
- Tensorflow version (GPU?):
- Using GPU in script?:
- Using distributed or parallel set-up in script?:


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