File tree 1 file changed +6
-6
lines changed
1 file changed +6
-6
lines changed Original file line number Diff line number Diff line change @@ -32,9 +32,9 @@ def test_compile_traced(self):
32
32
def test_compile_script (self ):
33
33
with torch .no_grad ():
34
34
trt_mod = torchtrt .ts .compile (self .scripted_model ,
35
- inputs = [self .input ],
36
- device = torchtrt .Device (gpu_id = 0 ),
37
- enabled_precisions = {torch .float })
35
+ inputs = [self .input ],
36
+ device = torchtrt .Device (gpu_id = 0 ),
37
+ enabled_precisions = {torch .float })
38
38
same = (trt_mod (self .input ) - self .scripted_model (self .input )).abs ().max ()
39
39
self .assertTrue (same < 2e-2 )
40
40
@@ -49,9 +49,9 @@ def test_compile_global(self):
49
49
def test_compile_global_nn_mod (self ):
50
50
with torch .no_grad ():
51
51
trt_mod = torchtrt .compile (self .model ,
52
- inputs = [self .input ],
53
- device = torchtrt .Device (gpu_id = 0 ),
54
- enabled_precisions = {torch .float })
52
+ inputs = [self .input ],
53
+ device = torchtrt .Device (gpu_id = 0 ),
54
+ enabled_precisions = {torch .float })
55
55
same = (trt_mod (self .input ) - self .scripted_model (self .input )).abs ().max ()
56
56
self .assertTrue (same < 2e-2 )
57
57
You can’t perform that action at this time.
0 commit comments