Add plugin arg to vmfbRunner

nod-ai · Jun 12, 2024 · fdaec4b · fdaec4b
1 parent 815c857
commit fdaec4b
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/models/turbine_models/model_runner.py b/models/turbine_models/model_runner.py
@@ -4,9 +4,13 @@
 
 
 class vmfbRunner:
-    def __init__(self, device, vmfb_path, external_weight_path=None):
+    def __init__(self, device, vmfb_path, external_weight_path=None, extra_plugin=None):
         flags = []
-        haldriver = ireert.get_driver(device)
+        clean_driver = False
+        if extra_plugin:
+            ireert.flags.parse_flags(f"--executable_plugin={extra_plugin}")
+            clean_driver = True
+        haldriver = ireert.get_driver(device, clean_driver)
         if "://" in device:
             try:
                 device_idx = int(device.split("://")[-1])