Merge pull request #81 from NVIDIA-AI-IOT/prepare_broadcastable_tensors

added mnasnet to test
NVIDIA-AI-IOT · Oct 1, 2019 · ad404c8 · ad404c8
2 parents b81b19f + 25c2e55
commit ad404c8
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),
 ]