Update app.py

app.py

@@ -17,7 +17,8 @@ tokenizer = AutoTokenizer.from_pretrained(model_path)
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda', quantization_config=quantization_config)
 rm_tokenizer = AutoTokenizer.from_pretrained('OpenAssistant/reward-model-deberta-v3-large-v2')
-rm_model = AutoModelForSequenceClassification.from_pretrained('OpenAssistant/reward-model-deberta-v3-large-v2', torch_dtype=torch.bfloat16)
+rm_model = AutoModelForSequenceClassification.from_pretrained('OpenAssistant/reward-model-deberta-v3-large-v2', device_map='cuda', torch_dtype=torch.bfloat16)
+
 
 @spaces.GPU
 def generate_text(usertitle, content, temperature, max_length, N=3):
@@ -30,6 +31,8 @@ def generate_text(usertitle, content, temperature, max_length, N=3):
     def score(sequence):
         inputs = rm_tokenizer(sequence, return_tensors='pt', padding=True, truncation=True, max_length=512).to('cuda')
         inputs = {k: v.to('cuda') for k, v in inputs.items()} 
+#       print(rm_model.device)
+#       print(inputs['input_ids'].device) 
         with torch.no_grad():
             outputs = rm_model(**inputs)
             logits = outputs.logits