resolve CUDA OOM issue in tests

zewenli98 · zewenli98 · commit a54907e103ae · 2025-12-04T10:33:02.000-08:00
diff --git a/tests/py/dynamo/models/test_weight_stripped_engine.py b/tests/py/dynamo/models/test_weight_stripped_engine.py
@@ -29,7 +29,7 @@ class TestWeightStrippedEngine(TestCase):
     )
     def test_three_ways_to_compile(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
-        example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
+        example_inputs = (torch.randn((2, 3, 224, 224)).to("cuda"),)
         exp_program = torch.export.export(pyt_model, example_inputs)
 
         settings = {

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ class TestWeightStrippedEngine(TestCase):`
`29`	`29`	`)`
`30`	`30`	`def test_three_ways_to_compile(self):`
`31`	`31`	`pyt_model = models.resnet18(pretrained=True).eval().to("cuda")`
`32`		`- example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)`
	`32`	`+ example_inputs = (torch.randn((2, 3, 224, 224)).to("cuda"),)`
`33`	`33`	`exp_program = torch.export.export(pyt_model, example_inputs)`
`34`	`34`
`35`	`35`	`settings = {`