fixed bug where device_idx was hardcoded (nod-ai#1693)

Co-authored-by: Elias Joseph <elias@nod-labs.com>
gpetters-amd · Jul 26, 2023 · 9d399eb · 9d399eb
1 parent 927b662
commit 9d399eb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/apps/language_models/scripts/vicuna.py b/apps/language_models/scripts/vicuna.py
@@ -785,7 +785,7 @@ def compile_to_vmfb_one_model(
                 module = SharkInference(
                     None,
                     device=device,
-                    device_idx=idx % 4,
+                    device_idx=device_idx,
                     mlir_dialect="tm_tensor",
                     mmap=False,
                 )
@@ -798,7 +798,7 @@ def compile_to_vmfb_one_model(
                 module = SharkInference(
                     mlirs[idx],
                     device=device,
-                    device_idx=idx % 4,
+                    device_idx=device_idx,
                     mlir_dialect="tm_tensor",
                     mmap=False,
                 )