From 97e721bbdd52ae5144fdea7dbad65275dba68cf7 Mon Sep 17 00:00:00 2001 From: Alagris Date: Sun, 23 May 2021 20:31:41 +0200 Subject: [PATCH] fix --- train_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_model.py b/train_model.py index fb441dd..a2e8049 100644 --- 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.numpy(), mock_y.numpy()) + return mlflow.models.signature.infer_signature(mock_text.detach().numpy(), mock_y.detach().numpy()) @ex.automain