Sending inputs to device for non- multi gpu hardware

src/deepeval/base_task.py

@@ -18,7 +18,18 @@ class BaseTask(ABC):
     def __init__(self, dataset_repo, model_name):
         self.dataset_repo = dataset_repo
         self.dataset = self.load_dataset_from_hf()
-        self.device = "auto" if torch.cuda.is_available() else "cpu"
+
+        device_count = torch.cuda.device_count()
+        if device_count > 1:
+            self.device = "auto"
+            print(f"Using {device_count} GPUs with auto config.")
+        elif device_count == 1:
+            self.device = "cuda"
+            print(f"Using {device_count} GPU with cuda config.")
+        else:
+            self.device = "cpu"
+            print("No GPU found. Using CPU.")
+
         self.model, self.tokenizer = self.get_cached_model(model_name, self.device)
         openai.api_key = OPENAI_KEY
 
@@ -130,8 +141,13 @@ class BaseTask(ABC):
         formatted_chat = self.tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
         #print(formatted_chat)
         inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
-        input_ids = inputs.input_ids
-        attention_mask = inputs.attention_mask
+
+        if self.device == "auto":
+            input_ids = inputs.input_ids
+            attention_mask = inputs.attention_mask
+        else:
+            input_ids = inputs.input_ids.to(self.model.device)
+            attention_mask = inputs.attention_mask.to(self.model.device)
 
         # Generate the sequence of letters starting from 'A'
         letters = [chr(ord('A') + i) for i in range(len(choices))]  # Create option letters A, B, C, D, E, ...
@@ -188,8 +204,13 @@ class BaseTask(ABC):
         )
 
         inputs = self.tokenizer(formatted_chat, return_tensors="pt", padding=True, truncation=True)
-        input_ids = inputs.input_ids
-        attention_mask = inputs.attention_mask
+        
+        if self.device == "auto":
+            input_ids = inputs.input_ids
+            attention_mask = inputs.attention_mask
+        else:
+            input_ids = inputs.input_ids.to(self.model.device)
+            attention_mask = inputs.attention_mask.to(self.model.device)
 
         output = self.model.generate(
             input_ids,

src/deepeval/instruction_following_task.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 from src.deepeval.base_task import BaseTask
 from deepeval.metrics import PromptAlignmentMetric
 from deepeval.test_case import LLMTestCase

src/deepeval/toxicity_task.py

@@ -1,4 +1,4 @@
-import datetime
+from datetime import datetime
 from src.deepeval.base_task import BaseTask
 from deepeval.metrics import ToxicityMetric
 from deepeval.test_case import LLMTestCase