From 7fefc558a0667f478df3f9494379ccc8728bff05 Mon Sep 17 00:00:00 2001
From: Vibhav Vineet <vibhav.vineet@gmail.com>
Date: Fri, 13 Dec 2024 15:58:36 -0800
Subject: [PATCH] add llava multigpu capability

---
 eureka_ml_insights/models/models.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/eureka_ml_insights/models/models.py b/eureka_ml_insights/models/models.py
index 5681171..d9e6334 100644
--- a/eureka_ml_insights/models/models.py
+++ b/eureka_ml_insights/models/models.py
@@ -648,7 +648,7 @@ def pick_available_device(self):
         import numpy as np
         import torch
 
-        device = "cpu"
+        device = "cpu"        
 
         if torch.cuda.is_available():
             utilizations = []
@@ -660,6 +660,9 @@ def pick_available_device(self):
 
             device = f"cuda:{gpu_index}"
 
+            if self.multi_gpu:
+                device = f"cuda"
+
         logging.info(f"Using device {device} for model self hosting")
 
         return device
@@ -738,6 +741,7 @@ class LLaVAHuggingFaceModel(HuggingFaceModel):
 
     quantize: bool = False
     use_flash_attn: bool = False
+    multi_gpu: bool = False
 
     def __post_init__(self):
         super().__post_init__()