Fix device bug

app.py

@@ -49,6 +49,8 @@ st.write('Loading the pretrained model ...')
 model_name = "CarolXia/pii-kd-deberta-v2"
 # config = PeftConfig.from_pretrained(model_name)
 model = DebertaV2ForTokenClassification.from_pretrained(model_name, token=st.secrets["HUGGINGFACE_TOKEN"])
+if torch.cuda.is_available():
+    model = model.to("cuda")
 # Try quantization instead
 # model = AutoModelForTokenClassification.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
 tokenizer = DebertaV2Tokenizer.from_pretrained("microsoft/mdeberta-v3-base", token=st.secrets["HUGGINGFACE_TOKEN"])
@@ -61,8 +63,7 @@ pytorch_total_params = sum(p.numel() for p in model.parameters())
 torch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
 print(f'total params: {pytorch_total_params}. tunable params: {torch_total_params}')
 
-if torch.cuda.is_available():
-    model = model.to("cuda")
+
 
 # Sample text containing PII/PHI entities
 text = """