From e394fdca519b7102ff29219edc538cc6b66aafc5 Mon Sep 17 00:00:00 2001 From: Witold Woch Date: Sun, 14 May 2023 23:23:37 +0200 Subject: [PATCH] Zaktualizuj 'train.py' --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 8061ce4..e07529b 100644 --- a/train.py +++ b/train.py @@ -102,7 +102,7 @@ with mlflow.start_run() as run: loss.backward() optimizer.step() - # Infer model signature to log it +torch.save(model,"classificationn_model.pt") signature = infer_signature(X_train.numpy(), model(X_train).detach().numpy()) input_example = {"input": X_train[0].numpy().tolist()}