Allow strategy in configs

datamol-io · Aug 16, 2023 · 81d1d9d · 81d1d9d
1 parent 4d472d7
commit 81d1d9d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/graphium/config/_loader.py b/graphium/config/_loader.py
@@ -372,7 +372,7 @@ def load_trainer(
     cfg_trainer = deepcopy(config["trainer"])
 
     # Define the IPU plugin if required
-    strategy = "auto"
+    strategy = config["trainer"]["trainer"].get("strategy", "auto")
     if accelerator_type == "ipu":
         ipu_opts, ipu_inference_opts = _get_ipu_opts(config)
 
@@ -385,6 +385,9 @@ def load_trainer(
             precision=config["trainer"]["trainer"].get("precision"),
         )
 
+        if strategy != "auto":
+            raise ValueError("IPUs selected, but strategy is not set to 'auto'")
+
         from lightning_graphcore import IPUStrategy
 
         strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)