Update app.py

app.py

@@ -62,14 +62,19 @@ tasks = {
 }
 
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:30'
-# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
+os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 os.environ['CUDA_CACHE_DISABLE'] = '1'
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the model and tokenizer globally
+# Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True)
 model = AutoModel.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True).to(device)
 
+# If multiple GPUs are available, enable DataParallel
+if torch.cuda.device_count() > 1:
+    print("Using", torch.cuda.device_count(), "GPUs")
+    model = DataParallel(model)
+
 # Embedding requests and response queues
 embedding_request_queue = queue.Queue()
 embedding_response_queue = queue.Queue()