@@ -96,42 +96,6 @@ def test_resnet18_cpu_offload(ir):
9696 torch ._dynamo .reset ()
9797
9898
99- @unittest .skipIf (
100- not importlib .util .find_spec ("torchvision" ), "torchvision not installed"
101- )
102- def test_resnet18_torch_exec_ops (ir ):
103- model = models .resnet18 (pretrained = True ).eval ().to ("cuda" )
104- input = torch .randn ((1 , 3 , 224 , 224 )).to ("cuda" )
105-
106- compile_spec = {
107- "inputs" : [
108- torchtrt .Input (
109- min_shape = (1 , 3 , 224 , 224 ),
110- opt_shape = (8 , 3 , 224 , 224 ),
111- max_shape = (16 , 3 , 224 , 224 ),
112- dtype = torch .float32 ,
113- )
114- ],
115- "ir" : ir ,
116- "enabled_precisions" : {torch .float32 , torch .float16 },
117- "min_block_size" : 1 ,
118- "output_format" : "exported_program" ,
119- "cache_built_engines" : True ,
120- "reuse_cached_engines" : True ,
121- "torch_executed_ops" : {torch .ops .aten .matmul , "torch.ops.aten.add" },
122- }
123-
124- trt_mod = torchtrt .compile (model , ** compile_spec )
125- cos_sim = cosine_similarity (model (input ), trt_mod (input ))
126- assertions .assertTrue (
127- cos_sim > COSINE_THRESHOLD ,
128- msg = f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: { cos_sim } Threshold: { COSINE_THRESHOLD } " ,
129- )
130-
131- # Clean up model env
132- torch ._dynamo .reset ()
133-
134-
13599@pytest .mark .unit
136100@pytest .mark .parametrize ("dtype" , [torch .float16 , torch .bfloat16 , torch .float32 ])
137101@unittest .skipIf (
0 commit comments