chore: Adding a nox session to download all the models for you

Running `nox --session download_test_models-3` will download all the models used in testing for you Signed-off-by: Naren Dasan <[email protected]> Signed-off-by: Naren Dasan <[email protected]>
pytorch · Mar 8, 2022 · 06baa81 · 06baa81
1 parent 11bcb98
commit 06baa81
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/noxfile.py b/noxfile.py
@@ -322,3 +322,8 @@ def l2_multi_gpu_tests(session):
 def l2_multi_gpu_tests_host_deps(session):
     """Makes sure that Torch-TensorRT can operate on multi-gpu systems using host dependencies"""
     run_l2_multi_gpu_tests(session, use_host_env=True)
+
+@nox.session(python=["3"], reuse_venv=True)
+def download_test_models(session):
+    """Grab all the models needed for testing"""
+    download_models(session, use_host_env=True)