fix(stablelm): running on GPU

Signed-off-by: aarnphm-ec2-dev <[email protected]>
bentoml · Jun 11, 2023 · a5efb7f · a5efb7f
1 parent 8762a56
commit a5efb7f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/openllm/models/stablelm/modeling_stablelm.py b/src/openllm/models/stablelm/modeling_stablelm.py
@@ -107,6 +107,9 @@ def generate(self, prompt: str, **attrs: t.Any) -> list[str]:
             "stopping_criteria": StoppingCriteriaList([StopOnTokens()]),
         }
 
+        if torch.cuda.is_available():
+            self.model.cuda()
+
         inputs = t.cast("torch.Tensor", self.tokenizer(prompt, return_tensors="pt")).to(self.device)
         tokens = self.model.generate(**inputs, **generation_kwargs)
         return [self.tokenizer.decode(tokens[0], skip_special_tokens=True)]