qwen to device

create_app.py

@@ -21,6 +21,7 @@ def load_models():
     global MODELS_LOADED, LONGFORMER_TOKENIZER, LONGFORMER_MODEL, QWEN_TOKENIZER, QWEN_MODEL
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if not MODELS_LOADED:
+        print("DEVICE", device)
         LONGFORMER_TOKENIZER = LongformerTokenizer.from_pretrained('allenai/longformer-base-4096', device='auto')
         config = LongformerConfig.from_json_file("Longformer_checkpoint/config.json")
         LONGFORMER_MODEL = CustomLongformerForSequenceClassification(config).from_pretrained('SFM2001/LongFormerScorer')
@@ -31,6 +32,7 @@ def load_models():
         QWEN_TOKENIZER = AutoTokenizer.from_pretrained(model_name, device='auto')
         QWEN_TOKENIZER.pad_token_id = QWEN_TOKENIZER.eos_token_id
         QWEN_MODEL = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16).half()
+        QWEN_MODEL = QWEN_MODEL.to(device)
         MODELS_LOADED = True
 
 def create_app():