File tree Expand file tree Collapse file tree 1 file changed +6
-5
lines changed
tests/py/dynamo/conversion Expand file tree Collapse file tree 1 file changed +6
-5
lines changed Original file line number Diff line number Diff line change 1
1
import torch
2
2
import torch .nn as nn
3
+ import torch_tensorrt
3
4
from parameterized import parameterized
4
5
from torch .testing ._internal .common_utils import TestCase , run_tests
5
- import torch_tensorrt
6
6
7
7
rand_ops = [
8
8
(
@@ -75,13 +75,14 @@ def forward(self):
75
75
return self .rand_op (self .size )
76
76
77
77
grid_model = TestModule (op , shape_or_input )
78
- #cannot use self.run_test() since it expects input in form of tensor
79
-
80
- #self.run_test(grid_model, None)
78
+ # cannot use self.run_test() since it expects input in form of tensor
79
+
80
+ # self.run_test(grid_model, None)
81
81
fx_graph = torch .fx .symbolic_trace (grid_model )
82
82
torch ._dynamo .reset ()
83
83
84
- optimized_model = torch_tensorrt .compile (fx_graph ,
84
+ optimized_model = torch_tensorrt .compile (
85
+ fx_graph ,
85
86
"torch_compile" ,
86
87
None ,
87
88
min_block_size = 1 ,
You can’t perform that action at this time.
0 commit comments