fix so that we check exact paths as well for is_prepared (#782)

nod-ai · Nov 6, 2024 · dddc7b1 · dddc7b1
1 parent cdc8703
commit dddc7b1
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/models/turbine_models/custom_models/pipeline_base.py b/models/turbine_models/custom_models/pipeline_base.py
@@ -540,7 +540,11 @@ def is_prepared(self, vmfbs, weights):
                 avail_files = os.listdir(self.external_weights_dir)
                 candidates = []
                 for filename in avail_files:
-                    if all(str(x) in filename for x in w_keywords):
+                    if all(
+                        str(x) in filename
+                        or str(x) == os.path.join(self.external_weights_dir, filename)
+                        for x in w_keywords
+                    ):
                         candidates.append(
                             os.path.join(self.external_weights_dir, filename)
                         )