Fix float32 indices in test_batched_export_with_backprop

jeanschmidt · web-flow · commit 78a66890e854 · 2026-04-21T04:57:53.000Z
Differential Revision: D101547370 Pull Request resolved: #18992
diff --git a/examples/models/llama/tests/test_static_attention.py b/examples/models/llama/tests/test_static_attention.py
@@ -325,7 +325,7 @@ def test_batched_export_with_backprop(self):
             static_config, input_len, cache_len, batch_size=batch_size
         )
         example_inputs = (
-            torch.zeros(batch_size, input_len),
+            torch.zeros(batch_size, input_len, dtype=torch.long),
             {
                 "masks": mgr.masks,
                 "freqs_cos_override": mgr.freqs_cos[:input_len],
@@ -350,7 +350,7 @@ def test_batched_export_with_backprop(self):
             static_config, input_len, cache_len, batch_size=1
         )
         example_inputs = (
-            torch.zeros(1, input_len),
+            torch.zeros(1, input_len, dtype=torch.long),
             {
                 "masks": mgr.masks,
                 "freqs_cos_override": mgr.freqs_cos[:input_len],

Original file line number	Diff line number	Diff line change
`@@ -325,7 +325,7 @@ def test_batched_export_with_backprop(self):`
`325`	`325`	`static_config, input_len, cache_len, batch_size=batch_size`
`326`	`326`	`)`
`327`	`327`	`example_inputs = (`
`328`		`- torch.zeros(batch_size, input_len),`
	`328`	`+ torch.zeros(batch_size, input_len, dtype=torch.long),`
`329`	`329`	`{`
`330`	`330`	`"masks": mgr.masks,`
`331`	`331`	`"freqs_cos_override": mgr.freqs_cos[:input_len],`
`@@ -350,7 +350,7 @@ def test_batched_export_with_backprop(self):`
`350`	`350`	`static_config, input_len, cache_len, batch_size=1`
`351`	`351`	`)`
`352`	`352`	`example_inputs = (`
`353`		`- torch.zeros(1, input_len),`
	`353`	`+ torch.zeros(1, input_len, dtype=torch.long),`
`354`	`354`	`{`
`355`	`355`	`"masks": mgr.masks,`
`356`	`356`	`"freqs_cos_override": mgr.freqs_cos[:input_len],`