From 1790ce99d2b15d540ed456baf78346dd3fdc7800 Mon Sep 17 00:00:00 2001
From: Jimmy Zhang <jiemingz@nvidia.com>
Date: Tue, 9 Apr 2024 15:53:22 -0700
Subject: [PATCH] gpu_per_node as arg

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
---
 nemo/export/tensorrt_llm.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nemo/export/tensorrt_llm.py b/nemo/export/tensorrt_llm.py
index b38eb592529c..d987eca2a7a8 100644
--- a/nemo/export/tensorrt_llm.py
+++ b/nemo/export/tensorrt_llm.py
@@ -245,6 +245,7 @@ def build(
         max_input_len: int = 256,
         max_output_len: int = 256,
         max_batch_size: int = 8,
+        gpus_per_node: int = 8,
         use_refit: bool = False,
         reshard_model: bool = False,
     ):