refactor: More useful torch version information

Signed-off-by: Naren Dasan <[email protected]>
pytorch · Jul 22, 2022 · 8fa9fbf · 8fa9fbf
1 parent 2be072a
commit 8fa9fbf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/py/requirements.txt b/py/requirements.txt
@@ -1,5 +1,3 @@
--f https://download.pytorch.org/whl/torch_stable.html
--f https://download.pytorch.org/whl/torch/
 --extra-index-url https://download.pytorch.org/whl/cu113
-torch==1.11.0+cu113
+torch==1.11.0
 pybind11==2.6.2
diff --git a/py/torch_tensorrt/_util.py b/py/torch_tensorrt/_util.py
@@ -1,6 +1,8 @@
 from torch_tensorrt import __version__
 from torch_tensorrt import _C
 
+import torch
+
 
 def dump_build_info():
     """Prints build information about the torch_tensorrt distribution to stdout
@@ -15,7 +17,9 @@ def get_build_info() -> str:
         str: String containing the build information for torch_tensorrt distribution
     """
     build_info = _C.get_build_info()
-    build_info = "Torch-TensorRT Version: " + str(__version__) + '\n' + build_info
+    build_info = "Torch-TensorRT Version: " + str(__version__) + '\n' \
+        + "Using PyTorch Version: " + str(torch.__version__) + '\n' \
+        + build_info
     return build_info