github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-08 05:24:00 -05:00
Code Issues Packages Projects Releases Wiki Activity

fixed bug where device_idx was hardcoded (#1693)

Browse Source 
Co-authored-by: Elias Joseph <elias@nod-labs.com>
This commit is contained in:
Eliasj42
2023-07-25 17:00:13 -07:00
committed by GitHub
parent 927b662aa7
commit 9d399eb988
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
apps/language_models/scripts/vicuna.py
View File

@@ -785,7 +785,7 @@ class ShardedVicuna(VicunaBase):
module = SharkInference(
None,
device=device,
device_idx=idx % 4,
device_idx=device_idx,
mlir_dialect="tm_tensor",
mmap=False,
)
@@ -798,7 +798,7 @@ class ShardedVicuna(VicunaBase):
module = SharkInference(
mlirs[idx],
device=device,
device_idx=idx % 4,
device_idx=device_idx,
mlir_dialect="tm_tensor",
mmap=False,
)