@@ -1281,7 +1281,12 @@ def forward(self, INPUT0, INPUT1):
12811281 )
12821282 return op0 , op1
12831283
1284- addSubModel = AddSubNet ((torch_output0_dtype , torch_output1_dtype , swap ,))
1284+ addSubModel = AddSubNet (
1285+ (
1286+ torch_output0_dtype ,
1287+ torch_output1_dtype , swap ,
1288+ )
1289+ )
12851290 traced = torch .jit .script (addSubModel )
12861291
12871292 model_version_dir = models_dir + "/" + model_name + "/" + str (model_version )
@@ -1345,7 +1350,11 @@ def forward(self, INPUT0, INPUT1):
13451350 return op0 , op1
13461351
13471352 ep = torch .export .export (
1348- AddSubNet (swap ), (torch .randn (* input_shape ), torch .randn (* input_shape ),),
1353+ AddSubNet (swap ),
1354+ (
1355+ torch .randn (* input_shape ),
1356+ torch .randn (* input_shape ),
1357+ ),
13491358 )
13501359 torch .export .save (ep , model_version_dir + "/model.pt2" )
13511360
@@ -1401,10 +1410,15 @@ def forward(self, INPUT0, INPUT1):
14011410 return op0 , op1
14021411
14031412 ep = torch .export .export (
1404- AddSubNet (swap ), (torch .randn (* input_shape ), torch .randn (* input_shape ),),
1413+ AddSubNet (swap ),
1414+ (
1415+ torch .randn (* input_shape ),
1416+ torch .randn (* input_shape ),
1417+ ),
14051418 )
14061419 torch ._inductor .aoti_compile_and_package (
1407- ep , package_path = model_version_dir + "/model.pt2" ,
1420+ ep ,
1421+ package_path = model_version_dir + "/model.pt2" ,
14081422 )
14091423
14101424
0 commit comments