Merge branch 'master' into wsad1-patch-1

pyg-team · Sep 30, 2024 · dd63def · dd63def
2 parents 700da76 + 6fc2761
commit dd63def
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 15 deletions.
diff --git a/test/nn/models/test_compile.py b/test/nn/models/test_compile.py
@@ -14,7 +14,7 @@
 from torch_frame.testing import withPackage
 
 
-@withPackage("torch>=2.1.0")
+@withPackage("torch>=2.5.0")
 @pytest.mark.parametrize(
     "model_cls, model_kwargs, stypes, expected_graph_breaks",
     [
@@ -34,7 +34,7 @@
                 gamma=0.1,
             ),
             None,
-            7,
+            2,
             id="TabNet",
         ),
         pytest.param(
@@ -47,21 +47,21 @@
                 ffn_dropout=0.5,
             ),
             None,
-            4,
+            0,
             id="TabTransformer",
         ),
         pytest.param(
             Trompt,
             dict(channels=8, num_prompts=2),
             None,
-            16,
+            4,
             id="Trompt",
         ),
         pytest.param(
             ExcelFormer,
             dict(in_channels=8, num_cols=3, num_heads=1),
             [stype.numerical],
-            4,
+            1,
             id="ExcelFormer",
         ),
     ],
@@ -89,4 +89,5 @@ def test_compile_graph_break(
         **model_kwargs,
     )
     explanation = torch._dynamo.explain(model)(tf)
-    assert explanation.graph_break_count <= expected_graph_breaks
+    graph_breaks = explanation.graph_break_count
+    assert graph_breaks == expected_graph_breaks
diff --git a/torch_frame/gbdt/tuned_lightgbm.py b/torch_frame/gbdt/tuned_lightgbm.py
@@ -103,7 +103,6 @@ def objective(
         trial: Any,  # optuna.trial.Trial
         train_data: Any,  # lightgbm.Dataset
         eval_data: Any,  # lightgbm.Dataset
-        cat_features: list[int],
         num_boost_round: int,
     ) -> float:
         r"""Objective function to be optimized.
@@ -112,8 +111,6 @@ def objective(
             trial (optuna.trial.Trial): Optuna trial object.
             train_data (lightgbm.Dataset): Train data.
             eval_data (lightgbm.Dataset): Validation data.
-            cat_features (list[int]): Array containing indexes of
-                categorical features.
             num_boost_round (int): Number of boosting round.
 
         Returns:
@@ -169,8 +166,7 @@ def objective(
 
         boost = lightgbm.train(
             self.params, train_data, num_boost_round=num_boost_round,
-            categorical_feature=cat_features, valid_sets=[eval_data],
-            callbacks=[
+            valid_sets=[eval_data], callbacks=[
                 lightgbm.early_stopping(stopping_rounds=50, verbose=False),
                 lightgbm.log_evaluation(period=2000)
             ])
@@ -199,19 +195,18 @@ def _tune(
         assert train_y is not None
         assert val_y is not None
         train_data = lightgbm.Dataset(train_x, label=train_y,
+                                      categorical_feature=cat_features,
                                       free_raw_data=False)
         eval_data = lightgbm.Dataset(val_x, label=val_y, free_raw_data=False)
 
         study.optimize(
             lambda trial: self.objective(trial, train_data, eval_data,
-                                         cat_features, num_boost_round),
-            num_trials)
+                                         num_boost_round), num_trials)
         self.params.update(study.best_params)
 
         self.model = lightgbm.train(
             self.params, train_data, num_boost_round=num_boost_round,
-            categorical_feature=cat_features, valid_sets=[eval_data],
-            callbacks=[
+            valid_sets=[eval_data], callbacks=[
                 lightgbm.early_stopping(stopping_rounds=50, verbose=False),
                 lightgbm.log_evaluation(period=2000)
             ])