Update multipurpose_chatbot/engines/transformers_engine.py

multipurpose_chatbot/engines/transformers_engine.py

@@ -4,6 +4,7 @@ import os
 import numpy as np
 import argparse
 import torch
+import sys
 import gradio as gr
 from typing import Any, Iterator
 from typing import Iterator, List, Optional, Tuple
@@ -427,6 +428,7 @@ class TransformersEngine(BaseEngine):
     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
+        import sys
         with torch.no_grad():
             inputs = self.tokenizer(prompt, return_tensors='pt')
             num_tokens = inputs.input_ids.size(1)