feat(onnx): use gpu

2024-06-19 18:28:15 +09:00
parent be7085f2f7
commit 1e6b04c0ec
3 changed files with 66 additions and 41 deletions
--- a/pyproject.toml
+++ b/pyproject.toml
@ -7,15 +7,16 @@ readme = "README.md"

 [tool.poetry.dependencies]
 python = "^3.10"
-onnxruntime = "^1.18.0"
 pillow = "^10.3.0"
 pandas = "^2.2.2"
 huggingface-hub = "^0.23.3"
 rich = "^13.7.1"
+onnxruntime-gpu = "^1.18.0"


 [tool.poetry.group.dev.dependencies]
 pytest = "^8.2.2"
+pytest-benchmark = "^4.0.0"

 [build-system]
 requires = ["poetry-core"]