fix

tasks/text.py

@@ -118,20 +118,20 @@ def bert_classifier(test_dataset: dict, model: str):
 
 def moe_classifier(test_dataset: dict, model: str):
     print("Starting MoE run")
+
+    # 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.to(device)
-    
+    embedding_model = AutoModel.from_pretrained("sentence-transformers/all-distilroberta-v1")    
     tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-distilroberta-v1")
 
     dataset = TextDataset(texts, tokenizer=tokenizer, max_length=512)
     dataloader = DataLoader(dataset, batch_size=64, shuffle=False)
 
-    # Use CUDA if available 
-    device, _, _ = get_backend()
-
     model = MoEClassifier(3, 0.05)
     model.load_state_dict(torch.load(model_path))
     model = model.to(device)