bugfix , support for torch 2.0 (#22)

* bug fix , support for torch 2.0 * add git ignore
ryujaehun · Jun 30, 2023 · 03b0dd5 · 03b0dd5
1 parent a60299c
commit 03b0dd5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+result/
diff --git a/benchmark_models.py b/benchmark_models.py
@@ -89,7 +89,8 @@ def train(precision="single"):
     benchmark = {}
     for model_type in MODEL_LIST.keys():
         for model_name in MODEL_LIST[model_type]:
-            model = getattr(model_type, model_name)(pretrained=False)
+            if model_name[-8:] == '_Weights': continue
+            model = getattr(model_type, model_name)()
             if args.NUM_GPU > 1:
                 model = nn.DataParallel(model, device_ids=range(args.NUM_GPU))
             model = getattr(model, precision)()
@@ -121,7 +122,8 @@ def inference(precision="float"):
     with torch.no_grad():
         for model_type in MODEL_LIST.keys():
             for model_name in MODEL_LIST[model_type]:
-                model = getattr(model_type, model_name)(pretrained=False)
+                if model_name[-8:] == '_Weights': continue
+                model = getattr(model_type, model_name)()
                 if args.NUM_GPU > 1:
                     model = nn.DataParallel(model, device_ids=range(args.NUM_GPU))
                 model = getattr(model, precision)()

diff --git a/requirement.txt b/requirement.txt
@@ -1,6 +1,6 @@
 matplotlib
 torchvision
-torch>=1.0.0
+torch==2.0.0
 pandas
 plotly
 cufflinks