Update inference_1.py

inference_1.py

@@ -17,7 +17,7 @@ def setup_model(model_path):
         model=model_path,
         trust_remote_code=True,
         # Optional parameters for performance tuning
-        # tensor_parallel_size=2,  # Use multiple GPUs
+        tensor_parallel_size=args.tensor_parallel_size,  # Use multiple GPUs
         # dtype="bfloat16",  # Use bfloat16 for more efficient inference
         # gpu_memory_utilization=0.85  # Control memory usage
     )
@@ -93,6 +93,7 @@ def main():
     parser.add_argument("--prompt", help="Prompt for single inference mode")
     parser.add_argument("--prompt_file", help="File containing prompts for batch mode (one per line)")
     parser.add_argument("--output_file", help="Path to save JSON results (default: auto-generated)")
+    parser.add_argument("--tensor_parallel_size", help="GPU number")
     args = parser.parse_args()
     
     # Initialize the model