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

Fix for CPU loading

Browse Source
This commit is contained in:
somebody
2023-06-21 21:18:43 -05:00
parent b81f61b820
commit 5ee20bd7d6
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
modeling/inference_models/hf_torch.py
View File

@@ -50,8 +50,12 @@ class BreakmodelConfig:
self.primary_device = 0 if torch.cuda.device_count() > 0 else "cpu"
def get_device_map(self, model: nn.Module) -> dict:
# HACK
if utils.args.cpu:
if (
# Explicitly CPU-only
utils.args.cpu
# No blocks are on GPU
or not sum(self.gpu_blocks)
):
self.primary_device = "cpu"
ram_blocks = len(utils.layers_module_names) - sum(self.gpu_blocks)