From 8fa9fbf538089d0c789b3d3b97a2a09ad1967d69 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Fri, 22 Jul 2022 12:06:19 -0700 Subject: [PATCH] refactor: More useful torch version information Signed-off-by: Naren Dasan --- py/requirements.txt | 4 +--- py/torch_tensorrt/_util.py | 6 +++++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/py/requirements.txt b/py/requirements.txt index 8d12c108aa..fce4b91dca 100644 --- 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 index ba260f8958..94a58cfcc5 100644 --- 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