gpu_per_node as arg

Signed-off-by: Jimmy Zhang <jiemingz@nvidia.com>
JimmyZhang12 · Apr 16, 2024 · 1790ce9 · 1790ce9
1 parent cd3c6e5
commit 1790ce9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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,
  ):