rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Actually get correct primary device

Browse Source
This commit is contained in:
somebody
2023-07-03 19:04:48 -05:00
parent 59c731f805
commit 6f7e6422ef
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
modeling/inference_models/hf_torch.py
View File

@@ -57,6 +57,10 @@ class BreakmodelConfig:
return "cpu"
elif torch.cuda.device_count() <= 0:
return "cpu"
for device_index, blocks in enumerate(self.gpu_blocks):
if blocks:
return device_index
return 0
def get_device_map(self, model: nn.Module) -> dict: