Update multipurpose_chatbot/engines/transformers_engine.py

multipurpose_chatbot/engines/transformers_engine.py

@@ -1,4 +1,5 @@
 
+import spaces
 import os
 import numpy as np
 import argparse
@@ -420,7 +421,8 @@ class TransformersEngine(BaseEngine):
         self._model.sample = types.MethodType(NewGenerationMixin.sample_stream, self._model)
         print(self._model)
         print(f"{self.max_position_embeddings=}")
-    
+
+    @spaces.GPU
     def generate_yield_string(self, prompt, temperature, max_tokens, stop_strings: Optional[Tuple[str]] = None, **kwargs):
         
         # ! MUST PUT INSIDE torch.no_grad() otherwise it will overflow OOM
@@ -428,7 +430,7 @@ class TransformersEngine(BaseEngine):
             inputs = self.tokenizer(prompt, return_tensors='pt')
             num_tokens = inputs.input_ids.size(1)
 
-            inputs = inputs.to(self.device_map)
+            inputs = inputs.to(self._model.device)
             
             generator = self._model.generate(
                 **inputs,