feat: Optimize hub.py download

Signed-off-by: Anurag Dixit <[email protected]>

Author: andi4191

.gitignore

@@ -57,4 +57,4 @@ examples/int8/qat/qat
 examples/int8/training/vgg16/data/*
 examples/int8/datasets/data/*
 env/**/*
-bazel-Torch-TensorRT-Preview
+model_snapshot.txt

tests/modules/hub.py

@@ -4,9 +4,30 @@
 import torchvision.models as models
 import timm
 from transformers import BertModel, BertTokenizer, BertConfig
+import os
+import sys
 
 torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
 
+torch_version = torch.__version__
+snapshot_file = 'model_snapshot.txt'
+skip_download = False
+
+# If model repository already setup
+if os.path.exists(snapshot_file):
+    with open(snapshot_file, 'r') as f:
+        model_version = f.read()
+        if model_version == torch_version:
+            skip_download = True
+
+# In case of existing model repository, skip the download
+if skip_download:
+    print('Skipping re-download of model repository')
+    sys.exit()
+else:
+    with open(snapshot_file, 'w') as f:
+        f.write(torch_version)
+
 models = {
     "alexnet": {
         "model": models.alexnet(pretrained=True),