add device

tasks/text.py

@@ -122,6 +122,8 @@ def moe_classifier(test_dataset: dict, model: str):
     model_path = f"tasks/text_models/0131_MoE_final.pt"
 
     embedding_model = AutoModel.from_pretrained("sentence-transformers/all-distilroberta-v1")
+    embedding_model.to(device)
+    
     tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-distilroberta-v1")
 
     dataset = TextDataset(texts, tokenizer=tokenizer, max_length=512)