-cpu flag

packages/node/meshnet_node/model_backend.py

@@ -212,6 +212,7 @@ class TorchModelShard:
         shard_end: int,
         quantization: Quantization = "auto",
         cache_dir: Path | None = None,
+        force_cpu: bool = False,
     ) -> None:
         if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
             raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -229,7 +230,10 @@ class TorchModelShard:
             ) from exc
 
         self.torch = torch
-        self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
+        if force_cpu:
+            self.device = torch.device("cpu")
+        else:
+            self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
         load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
         quant_config, dtype, uses_quantized_weights = _model_load_plan(
             AutoConfig,
@@ -682,8 +686,11 @@ def load_torch_shard(
     shard_end: int,
     quantization: Quantization = "auto",
     cache_dir: Path | None = None,
+    force_cpu: bool = False,
 ) -> TorchModelShard:
-    return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
+    return TorchModelShard(
+        model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu
+    )
 
 
 def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None: