From 85bceeb57d8c596e14398fb723ac57f0a5b8468d Mon Sep 17 00:00:00 2001
From: perib <pedroribeiroh@gmail.com>
Date: Mon, 22 Apr 2024 13:09:16 -0700
Subject: [PATCH] add classes_ to PytorchLRClassifier

---
 tpot/builtins/nn.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tpot/builtins/nn.py b/tpot/builtins/nn.py
index 953a9c0f..5c8dd049 100644
--- a/tpot/builtins/nn.py
+++ b/tpot/builtins/nn.py
@@ -120,6 +120,7 @@ def fit(self, X, y):
         # pylint: disable=no-member
 
         self._init_model(X, y)
+        self.classes_ = np.unique(y)
 
         assert _pytorch_model_is_fully_initialized(self)