Update run.py

run.py

@@ -33,8 +33,9 @@ def tp_tf_test(model_selector, queries_selector, prompt_selector, metric_selecto
 		return tokenize
 	
 	model = AutoModelForSequenceClassification.from_pretrained(model_selector)
-	device = torch.device("cuda")
-	model.to(device)
+	if torch.cuda.is_available():
+		device = torch.device("cuda")
+		model.to(device)
 	tokenizer = AutoTokenizer.from_pretrained(model_selector)	
 	data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 	
@@ -56,7 +57,10 @@ def tp_tf_test(model_selector, queries_selector, prompt_selector, metric_selecto
 		labels = []
 		nli_labels =[]
 		for batch in dataloader:
-			data = {k: v.to(device) for k, v in batch.items() if k not in ['labels', 'nli_label']}
+			if torch.cuda.is_available():
+				data = {k: v.to(device) for k, v in batch.items() if k not in ['labels', 'nli_label']}
+			else:
+				data = {k: v for k, v in batch.items() if k not in ['labels', 'nli_label']}
 			with torch.no_grad():
 				outputs = model(**data)
 			logits = outputs.logits			
@@ -100,18 +104,4 @@ def tp_tf_test(model_selector, queries_selector, prompt_selector, metric_selecto
 		for result in results:
 			results[result].to_csv(results_file, mode='a', index_label = result)
 			print(results[result], '\n')
-	return results
-
-
-
-
-
-
-
-
-
-
-
-
-
-
+	return results