InvalidArgumentError: Data t InvalidArgumentError: Data type mismatch at component 0: expected int32 but got int64 - Tensorflow

I am trying to train a model with Tensorflow following a course and I get the mentioned error.

here is the relevant part of my code:

element_spec = ({'input_ids': tf.TensorSpec(shape=(16, 512), dtype=tf.int32, name=None),
  'attention_masks': tf.TensorSpec(shape=(16, 512), dtype=tf.int32, name=None)},
 tf.TensorSpec(shape=(16, 5), dtype=tf.float64, name=None))

train_ds = tf.data.experimental.load('train', element_spec)
val_ds = tf.data.experimental.load('val', element_spec)

#in order to keep the history of our runs
history = model.fit(
    train_ds,
    validation_data = val_ds,
    epochs = 3
)

I tried running it on my laptop(no specific GPUs) and it is taking forever, so I decided to run the same thing on google colab but I get this error which I find odd:

Epoch 1/3

---------------------------------------------------------------------------

InvalidArgumentError                      Traceback (most recent call last)

<ipython-input-146-9b8dacb137ae> in <module>()
      3     train_ds,
      4     validation_data = val_ds,
----> 5     epochs = 3
      6 )

1 frames

/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     57     ctx.ensure_initialized()
     58     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
---> 59                                         inputs, attrs, num_outputs)
     60   except core._NotOkStatusException as e:
     61     if name is not None:

InvalidArgumentError:  Data type mismatch at component 0: expected int32 but got int64.
     [[node IteratorGetNext
 (defined at /usr/local/lib/python3.7/dist-packages/keras/engine/training.py:866)
]] [Op:__inference_train_function_16829]

Errors may have originated from an input operation.
Input Source operations connected to node IteratorGetNext:
In[0] iterator (defined at /usr/local/lib/python3.7/dist-packages/keras/engine/training.py:1216)

is the problem just somehow changing epoch = 3 from int64 to int32?

Solution 1:

Ok, I figured it out. I had defined dtype=tf.int32 in element spec. After changing it to dtype=tf.int64 it is working now.