fix device

tasks/text.py

@@ -121,12 +121,13 @@ def moe_classifier(test_dataset: dict, model: str):
 
     # Use CUDA if available 
     device, _, _ = get_backend()
-    
+
     texts = test_dataset["quote"]
 
     model_path = f"tasks/text_models/0131_MoE_final.pt"
 
-    embedding_model = AutoModel.from_pretrained("sentence-transformers/all-distilroberta-v1")    
+    embedding_model = AutoModel.from_pretrained("sentence-transformers/all-distilroberta-v1")
+    embedding_model = embedding_model.to(device)
     tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-distilroberta-v1")
 
     dataset = TextDataset(texts, tokenizer=tokenizer, max_length=512)