github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-09 22:07:55 -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( module = SharkInference(
None, None,
device=device, device=device,
device_idx=idx % 4, device_idx=device_idx,
mlir_dialect="tm_tensor", mlir_dialect="tm_tensor",
mmap=False, mmap=False,
) )
@@ -798,7 +798,7 @@ class ShardedVicuna(VicunaBase):
module = SharkInference( module = SharkInference(
mlirs[idx], mlirs[idx],
device=device, device=device,
device_idx=idx % 4, device_idx=device_idx,
mlir_dialect="tm_tensor", mlir_dialect="tm_tensor",
mmap=False, mmap=False,
) )