fix bug

chengtao-lv · chengtao-lv · commit 121c6b125f6c · 2025-07-15T13:20:49.000+08:00
diff --git a/llmc/compression/token_reduction/divprune.py b/llmc/compression/token_reduction/divprune.py
@@ -129,7 +129,7 @@ def wrapper(self, *args, **kwargs):
                 token_indices = input_ids[0][attention_mask[0]] == IMAGE_TOKEN_INDEX
                 pruning_paras['image_token_start_index'] = torch.where(token_indices)[
                     0
-                ].item()
+                ][0].item()
 
                 outputs = fn(*args, **kwargs)
 

Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ def wrapper(self, args, *kwargs):`
`129`	`129`	`token_indices = input_ids[0][attention_mask[0]] == IMAGE_TOKEN_INDEX`
`130`	`130`	`pruning_paras['image_token_start_index'] = torch.where(token_indices)[`
`131`	`131`	`0`
`132`		`- ].item()`
	`132`	`+ ][0].item()`
`133`	`133`
`134`	`134`	`outputs = fn(args, *kwargs)`
`135`	`135`