From 498442f705972b56109ccf1194fb7bd9064df6a0 Mon Sep 17 00:00:00 2001
From: Harutaka Kawamura <hkawamura0130@gmail.com>
Date: Thu, 25 Feb 2021 18:49:41 +0900
Subject: [PATCH] Replace IrisClassificationWithoutValidation with
 IrisClassification (#4141)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
---
 tests/pytorch/test_pytorch_autolog.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/pytorch/test_pytorch_autolog.py b/tests/pytorch/test_pytorch_autolog.py
index 5a71aa2c8aa74..40e90553f7e67 100644
--- a/tests/pytorch/test_pytorch_autolog.py
+++ b/tests/pytorch/test_pytorch_autolog.py
@@ -48,7 +48,7 @@ def pytorch_model_without_validation():
 @pytest.mark.parametrize("log_models", [True, False])
 def test_pytorch_autolog_log_models_configuration(log_models):
     mlflow.pytorch.autolog(log_models=log_models)
-    model = IrisClassificationWithoutValidation()
+    model = IrisClassification()
     dm = IrisDataModule()
     dm.prepare_data()
     dm.setup(stage="fit")