added mnasnet to test

NVIDIA-AI-IOT · Oct 1, 2019 · 25c2e55 · 25c2e55
1 parent 45cb4e2
commit 25c2e55
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/torch2trt/module_test.py b/torch2trt/module_test.py
@@ -40,6 +40,10 @@ def module_name(self):
     ModuleTest(torchvision.models.shufflenet_v2_x1_0, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
     ModuleTest(torchvision.models.shufflenet_v2_x1_5, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
     ModuleTest(torchvision.models.shufflenet_v2_x2_0, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
+    ModuleTest(torchvision.models.mnasnet0_5, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
+    ModuleTest(torchvision.models.mnasnet0_75, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
+    ModuleTest(torchvision.models.mnasnet1_0, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
+    ModuleTest(torchvision.models.mnasnet1_3, torch.float16, torch.device('cuda'), [(1, 3, 224, 224)], fp16_mode=True),
 ]