fix: move to self device (#309)

* move to self device * chlog Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Lightning-Universe · Nov 21, 2022 · 375be3e · 375be3e
1 parent 2e38190
commit 375be3e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed passing config name to `CNNDailyMailSummarizationDataModule` ([#310](https://github.com/Lightning-AI/lightning-transformers/pull/310))
 
 
+- Fixed move pipeline to `self.device` as default ([#309](https://github.com/Lightning-AI/lightning-transformers/pull/309))
+
+
 ## [0.2.4] - 2022-11-03
 
 ### Changed

diff --git a/lightning_transformers/core/model.py b/lightning_transformers/core/model.py
@@ -151,6 +151,7 @@ def hf_pipeline_task(self) -> Optional[str]:
     def hf_pipeline(self) -> Pipeline:
         if self._hf_pipeline is None:
             if self.hf_pipeline_task is not None:
+                self._hf_pipeline_kwargs["device"] = self._hf_pipeline_kwargs.get("device", self.device)
                 self._hf_pipeline = hf_transformers_pipeline(
                     task=self.hf_pipeline_task, model=self.model, tokenizer=self.tokenizer, **self._hf_pipeline_kwargs
                 )