fix libtorch pt2 generation

whoisj · whoisj · commit b94a33c3f148 · 2025-12-22T13:21:10.000-05:00
diff --git a/qa/common/gen_qa_models.py b/qa/common/gen_qa_models.py
@@ -1352,8 +1352,8 @@ def forward(self, INPUT0, INPUT1):
     ep = torch.export.export(
         AddSubNet(swap),
         (
-            torch.randn(*input_shape),
-            torch.randn(*input_shape),
+            torch.randn(*input_shape, device="cuda"),
+            torch.randn(*input_shape, device="cuda"),
         ),
     )
     torch.export.save(ep, model_version_dir + "/model.pt2")
@@ -2112,7 +2112,7 @@ def create_models(
         )
         create_libtorch_pt2_modelfile(
             models_dir,
-            8,
+            0,
             model_version,
             input_shape,
             output0_shape,

Original file line number	Diff line number	Diff line change
`@@ -1352,8 +1352,8 @@ def forward(self, INPUT0, INPUT1):`
`1352`	`1352`	`ep = torch.export.export(`
`1353`	`1353`	`AddSubNet(swap),`
`1354`	`1354`	`(`
`1355`		`- torch.randn(*input_shape),`
`1356`		`- torch.randn(*input_shape),`
	`1355`	`+ torch.randn(*input_shape, device="cuda"),`
	`1356`	`+ torch.randn(*input_shape, device="cuda"),`
`1357`	`1357`	`),`
`1358`	`1358`	`)`
`1359`	`1359`	`torch.export.save(ep, model_version_dir + "/model.pt2")`
`@@ -2112,7 +2112,7 @@ def create_models(`
`2112`	`2112`	`)`
`2113`	`2113`	`create_libtorch_pt2_modelfile(`
`2114`	`2114`	`models_dir,`
`2115`		`- 8,`
	`2115`	`+ 0,`
`2116`	`2116`	`model_version,`
`2117`	`2117`	`input_shape,`
`2118`	`2118`	`output0_shape,`