Update model.py (#739)

* Update model.py * Fix review comment
2026-04-20 03:02:16 -04:00 · 2022-10-02 05:55:53 +05:30
parent 74557472ec
commit 0df54ec3f3
1 changed files with 1 additions and 1 deletions
--- a/flaml/model.py
+++ b/flaml/model.py
@@ -2287,7 +2287,7 @@ class TemporalFusionTransformerEstimator(SKLearnEstimator):
            kwargs.get("log_dir", "lightning_logs")
        )  # logging results to a tensorboard
        default_trainer_kwargs = dict(
-            gpus=self._kwargs.get("gpu_per_trial", [0])
+            gpus=kwargs.get("gpu_per_trial", [0])
            if torch.cuda.is_available()
            else None,
            max_epochs=max_epochs,