try numba

requirements.txt

@@ -3,4 +3,6 @@ keras==2.3.1
 pandas==1.0.5
 scikit-learn==0.23.1
 
+numba==0.50.1
+
 redis==2.10.5

src/models/esmm/model.py

 import timeit
 
+import numba
 import tensorflow as tf
 from tensorflow import feature_column as fc
 from tensorflow.python.estimator.canned import head as head_lib
@@ -91,6 +92,7 @@ def _bytes_feature(value):
     return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
 
+@numba.jit(nopython=True, parallel=True)
 def model_predict(inputs, predict_fn):
     time_1 = timeit.default_timer()
     int_columns = [