mask decoder_input_ids

Signed-off-by: Kyle Sayers <[email protected]>

src/llmcompressor/modifiers/utils/pytorch_helpers.py

@@ -41,8 +41,14 @@ def apply_pad_mask_to_batch(batch: Dict[str, torch.Tensor]) -> Dict[str, torch.T
     :param batch: batch to apply padding to if it exists
     :return: batch with padding zeroed out in the input_ids
     """
-    if "input_ids" in batch and "attention_mask" in batch:
-        batch["input_ids"] = batch["input_ids"] * batch["attention_mask"]
+    if "attention_mask" in batch:
+        if "input_ids" in batch:
+            batch["input_ids"] = batch["input_ids"] * batch["attention_mask"]
+        if "decoder_input_ids" in batch:
+            batch["decoder_input_ids"] = (
+                batch["decoder_input_ids"] * batch["attention_mask"]
+            )
+
     return batch