The solution is to manually set the cardinality as below:
# print(len(train_data)) gives error
train_data = train_data.apply(tf.data.experimental.assert_cardinality(NUM_BATCHES))
print(len(train_data)) # NUM_BATCHES