Skip to content

Commit

Permalink
Allow plugin benchmark tools to override the built-in one
Browse files Browse the repository at this point in the history
  • Loading branch information
jeremyfowers committed Aug 12, 2024
1 parent b0acc01 commit 779536f
Showing 1 changed file with 8 additions and 2 deletions.
10 changes: 8 additions & 2 deletions src/turnkeyml/sequence/tool_plugins.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import turnkeyml.tools.onnx as onnx_tools
import turnkeyml.common.plugins as plugins
import turnkeyml.tools.management_tools as mgmt
from turnkeyml.run.benchmark_model import Benchmark
from turnkeyml.tools.discovery import Discover
import turnkeyml.tools.report as report
from turnkeyml.tools.load_build import LoadBuild
Expand All @@ -16,7 +15,6 @@
mgmt.Cache,
mgmt.ModelsLocation,
report.Report,
Benchmark,
Discover,
export.ExportPytorchModel,
onnx_tools.OptimizeOnnxModel,
Expand All @@ -40,3 +38,11 @@
)

SUPPORTED_TOOLS.append(tool_class)

# Give a "benchmark" tool installed by a plugin priority over
# a "benchmark" tool built into turnkeyml
tool_names = [tool.unique_name for tool in SUPPORTED_TOOLS]
if "benchmark" not in tool_names:
from turnkeyml.run.benchmark_model import Benchmark

SUPPORTED_TOOLS.append(Benchmark)

0 comments on commit 779536f

Please sign in to comment.