title: TensorFlow ResourceExhaustedError when running a simple model on Google Colab
body: <p>I'm trying to understand why I'm getting a ResourceExhaustedError when running a simple model:</p>
  <pre class="lang-py prettyprint-override"><code>model = models.Sequential([
    layers.Input(shape=(None, 1)),
    layers.Conv1D(padding='same', kernel_initializer='Orthogonal', filters=16, kernel_size=16, activation=None, strides=2),
    layers.Dropout(0.5),
    layers.Dense(units=80000),
  ])
  
  model.compile(loss='mean_squared_error', optimizer='sgd', metrics=['accuracy'])
  
  history = model.fit(train_waveforms_ds,
                      validation_data=validation_waveforms_ds,
                      epochs=100)
  
  </code></pre>
  <p>The error I get:</p>
  <pre><code>_________________________________________________________________
  None
  Epoch 1/50
  ---------------------------------------------------------------------------
  ResourceExhaustedError                    Traceback (most recent call last)
  &lt;ipython-input-67-fae874124f81&gt; in &lt;module&gt;()
       12 history = model.fit(train_waveforms_ds,
       13                     validation_data=validation_waveforms_ds,
  ---&gt; 14                     epochs=EPOCHS)
  
  6 frames
  /usr/local/lib/python3.6/dist-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,
  ---&gt; 60                                         inputs, attrs, num_outputs)
       61   except core._NotOkStatusException as e:
       62     if name is not None:
  
  ResourceExhaustedError:  OOM when allocating tensor with shape[80000,80000] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
       [[node sequential_34/dense_34/Tensordot/MatMul (defined at &lt;ipython-input-67-fae874124f81&gt;:14) ]]
  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.
   [Op:__inference_train_function_16700]
  
  Function call stack:
  train_function
  </code></pre>
  <p>Here's a print of <code>model.summary()</code></p>
  <pre><code>Model: &quot;sequential_34&quot;
  _________________________________________________________________
  Layer (type)                 Output Shape              Param #   
  =================================================================
  conv1d_37 (Conv1D)           (None, None, 16)          272       
  _________________________________________________________________
  dropout_15 (Dropout)         (None, None, 16)          0         
  _________________________________________________________________
  dense_34 (Dense)             (None, None, 80000)       1360000   
  =================================================================
  Total params: 1,360,272
  Trainable params: 1,360,272
  Non-trainable params: 0
  
  </code></pre>
  <p>I'm running this on a Google Colab using the GPU runtime. When I run the same code in a Gooogle Colab using the TPU runtime I just get a crash.</p>
  <p>I'm feeding <code>tf.data.Dataset</code>s into the model, so I don't think I can control the batch size...</p>
  <p>Elements of a waveform dataset's x have shape (None, None, 40000), and elements of a waveform dataset's y have shape (None, None, 80000).</p>