update torch dep to ^2.0 for consistency with ml4gw

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

Hide file tree

Showing 3 changed files with 279 additions and 26 deletions.
diff --git a/hermes/quiver/exporters/torch_onnx.py b/hermes/quiver/exporters/torch_onnx.py
@@ -14,7 +14,7 @@
 from hermes.quiver.exporters import Exporter
 
 
-def get_input_names_from_script_module(m: torch.jit.ScriptModule):
+def get_input_names_from_script_module(m):
     graph = m.graph
     input_names = [node.debugName().split(".")[0] for node in graph.inputs()]
     if "self" in input_names:

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -30,7 +30,7 @@ requests = "^2.26.0"
 
 # quiver optional dependencies
 tensorflow = {version = "^2.3", optional = true}
-torch = {version = "^1.7", 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" }