fix

2021-05-23 20:19:22 +02:00 · 2021-05-23 20:19:22 +02:00 · 1dda33aec7
commit 1dda33aec7
parent b8340d0604
1 changed files with 1 additions and 1 deletions
--- a/train_model.py
+++ b/train_model.py
@ -197,7 +197,7 @@ def signature(model, in_alphabet, max_len):
    mock_x = [('abc', 'xyz')]
    mock_text, _ = encode_str(mock_x, in_alphabet, max_len)
    mock_y = model(mock_text)
-    return mlflow.models.signature.infer_signature(mock_text, mock_y)
+    return mlflow.models.signature.infer_signature(mock_text.numpy(), mock_y.numpy())


@ex.automain