TensorFlow savedmodel input dimension does not match

[ZhanqiuHu]

When I used the saved TensorFlow query tower model created using offline-batch-recsys/Offline-Batch-Recommender-System.ipynb for Triton inference server, I got an error saying

1-user-embeddings      UNAVAILABLE: Invalid argument: model '1-user-embeddings', tensor 'user_age': the model expects 1 dimensions ( shape [-1]) but the model configuration specifies 2 dimensions (shape [-1,1])

I'm wondering if there is a way to create the query tower model so that the input shape becomes (shape [-1,1]) instead of shape [-1].

Thanks!