Update code/inference.py

code/inference.py

@@ -12,11 +12,12 @@ def model_fn(model_dir):
         model=model_dir,
         trust_remote_code=True,
         dtype="float16",
+        tensor_parallel_size=4,  # Use 4 GPUs for parallelization
         gpu_memory_utilization=0.9,
     )
     return model
 
-def predict_fn(data, model,context = None):
+def predict_fn(data, model, context=None):
     try:
         input_text = data.pop("inputs", data)
         parameters = data.pop("parameters", {})
@@ -36,11 +37,11 @@ def predict_fn(data, model,context = None):
         logger.error(f"Exception during prediction: {e}")
         return {"error": str(e)}
 
-def input_fn(request_body, request_content_type, context = None):
+def input_fn(request_body, request_content_type, context=None):
     if request_content_type == "application/json":
         return json.loads(request_body)
     else:
         raise ValueError(f"Unsupported content type: {request_content_type}")
 
-def output_fn(prediction, accept, context = None):
-    return json.dumps(prediction)
+def output_fn(prediction, accept, context=None):
+    return json.dumps(prediction)