Update miniscene2behavior.py

Imageomics · Dec 21, 2024 · 7e6cdc7 · 7e6cdc7
1 parent 232f85b
commit 7e6cdc7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/kabr_tools/miniscene2behavior.py b/src/kabr_tools/miniscene2behavior.py
@@ -69,6 +69,10 @@ def create_model(config_path: str, checkpoint_path: str, gpu_num: int) -> tuple[
     assert checkpoint_path is not None
     assert gpu_num >= 0
 
+    # load config
+    cfg = load_config(config_path)
+    cfg.NUM_GPUS = gpu_num
+
     # set random seed
     random.seed(cfg.RNG_SEED)
     np.random.seed(cfg.RNG_SEED)