Update test_latency.py

fix import issue, use deploy state during speed test
huawei-noah · May 30, 2023 · 37ceed6 · 37ceed6
1 parent f5d8ec9
commit 37ceed6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test_latency.py b/test_latency.py
@@ -9,7 +9,7 @@
 import torchvision
 import time
 
-import models.vanillanet
+from models.vanillanet import *
 
 
 if __name__ == "__main__":
@@ -22,6 +22,7 @@
 
     net = vanillanet_5().cuda()
     net.eval()
+    net.switch_to_deploy()
     print(net)
     for img, target in data_loader_val:
         img = img.cuda()