add onnx dep

ML4GW · Feb 6, 2024 · de890ac · de890ac
1 parent bfdcf47
commit de890ac
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 2 deletions.
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -33,11 +33,12 @@ tensorflow = {version = "^2.3", optional = true}
 torch = {version = "^2.0", optional = true}
 google-cloud-storage = {version = "^1.38", optional = true }
 nvidia-tensorrt = { version = "^8.0", optional = true, source = "ngc" }
+onnx = {version = "^1.15.0", optional = true}
 
 [tool.poetry.extras]
 # quiver 
 tensorflow = ["tensorflow"]
-torch = ["torch", "urllib3"]
+torch = ["torch", "onnx", "urllib3"]
 gcs = ["google-cloud-storage"]
 tensorrt = ["nvidia-tensorrt"]