Skip to content

Commit a0a3c43

Browse files
committed
feat: Add device mapping for llmlingua2
1 parent c32ee0b commit a0a3c43

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

huginn_hears/main.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,8 @@ def load_extractive_summarizer(self):
113113
"""
114114
model = PromptCompressor(
115115
model_name=self.model_name,
116-
use_llmlingua2=True
116+
use_llmlingua2=True,
117+
device_map="cuda" if torch.cuda.is_available() else "cpu"
117118
)
118119
try:
119120
yield model

0 commit comments

Comments
 (0)