fix a bug in the definition of source data shape (#335)

microsoft · Sep 25, 2024 · d902220 · d902220
1 parent 3009889
commit d902220
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 77 deletions.
diff --git a/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/playground-series-s4e8_template/model/model_nn.py
diff --git a/rdagent/scenarios/kaggle/experiment/scenario.py b/rdagent/scenarios/kaggle/experiment/scenario.py
@@ -24,6 +24,7 @@ def __init__(self, competition: str) -> None:
         super().__init__()
         self.competition = competition
         self.competition_descriptions = crawl_descriptions(competition)
+        self.input_shape = None
         self._source_data = self.source_data
 
         self.competition_type = None
@@ -121,6 +122,7 @@ def source_data(self) -> str:
             buffer = io.StringIO()
             X_valid.info(verbose=True, buf=buffer, show_counts=True)
             data_info = buffer.getvalue()
+            self.input_shape = X_valid.shape
             return data_info
 
         preprocess_experiment = KGFactorExperiment([])
@@ -141,7 +143,7 @@ def source_data(self) -> str:
         pickle.dump(X_test, open(data_folder / "X_test.pkl", "wb"))
         pickle.dump(others, open(data_folder / "others.pkl", "wb"))
 
-        self.input_shape = X_train.shape
+        self.input_shape = X_valid.shape
 
         buffer = io.StringIO()
         X_valid.info(verbose=True, buf=buffer, show_counts=True)