title: Model FIT using Keras works well only on CPU, on GPU loss NaN
body: <p>I have prepared this code, for image colorization:</p>
  <pre><code>train = dataGenerator.flow_from_directory(&quot;simple_images/data/&quot;, target_size=(256, 256), batch_size=340, class_mode=None)
  for img in train[0]:
    try:
        lab = rgb2lab(img)
        X.append(lab[:,:,0]) 
        Y.append(lab[:,:,1:] / 128)
    except:
       print('error')
  L = np.array(L)
  ab = np.array(ab)
  L = L.reshape(L.shape+(1,))
  </code></pre>
  <p>And then this model:</p>
  <pre><code>#Encoder
  model = Sequential()
  model.add(Conv2D(32, (3, 3), activation='relu', padding='same', input_shape=(256, 256, 1)))
  model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
  model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
  model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
  model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
  model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
  model.add(Conv2D(256, (3,3), activation='relu', padding='same'))
  #Decoder
  model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
  model.add(UpSampling2D((2, 2)))
  model.add(Conv2D(64, (3,3), activation='relu', padding='same'))
  model.add(UpSampling2D((2, 2)))
  model.add(Conv2D(32, (3,3), activation='relu', padding='same'))
  model.add(Conv2D(16, (3,3), activation='relu', padding='same'))
  model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
  model.add(UpSampling2D((2, 2)))
  model.compile(optimizer='adam', loss='mse' , metrics=['accuracy'])
  model.summary()
  history = model.fit(L,ab,validation_split=0.1, epochs=300, batch_size=64)
  </code></pre>
  <p>And when i run it, on CPU it works perfect, but it will took so much time, to train. So instead of CPU, i want to use my GPU RTX 3070. But when i run same code, on GPU i have got this:</p>
  <pre><code>Epoch 1/2
   2/20 [==&gt;...........................] - ETA: 1s - loss: nan - accuracy: 0.3925   WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0089s vs `on_train_batch_end` time: 0.0678s). Check your callbacks.
  20/20 [==============================] - 2s 78ms/step - loss: nan - accuracy: 0.3167 - val_loss: nan - val_accuracy: 0.2474
  Epoch 2/2
  20/20 [==============================] - 1s 69ms/step - loss: nan - accuracy: 0.3106 - val_loss: nan - val_accuracy: 0.2474
  INFO:tensorflow:Assets written to: colorize_autoencoder.model\assets
  </code></pre>
  <p>For demonstration i used only two epochs. How to fix it?</p>