bugfix v2: naming

Signed-off-by: Piotr Kaminski <pikaminski@nvidia.com>
NVIDIA · Aug 21, 2024 · 250525e · 250525e
1 parent 84a5e5e
commit 250525e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/nemo/export/trt_llm/converter/utils.py b/nemo/export/trt_llm/converter/utils.py
@@ -91,8 +91,9 @@ def save_split(split_vals, dir, key, i, split_factor):
 
 def save_expert_split(split_vals, dir, key, i, split_factor):
     for j, val in enumerate(split_vals):
-        suffix = f".{tp_num}.bin" if tp_num else ''
         tp_num = i * split_factor + j
+        suffix = f".{tp_num}.bin" if tp_num else ''
+
         global weights_dict
         weights_dict[key + suffix] = val