diff --git a/generator/ctrl/training_utils/training.py b/generator/ctrl/training_utils/training.py
index e8da996..dd94e8c 100644
--- a/generator/ctrl/training_utils/training.py
+++ b/generator/ctrl/training_utils/training.py
@@ -102,7 +102,7 @@ with tf.device('/gpu:0'):
 
 
     # input for the keras model
-    tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+    tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='float32')
 
 
     tied_embedding_softmax = TiedEmbeddingSoftmax()