Enable gradient checkpointing

2025-05-17 20:40:21 +08:00
parent 330cc0eeb4
commit 3d6fe064f5
3 changed files with 5 additions and 20 deletions
@@ -1030,8 +1030,8 @@ class ConformerEncoder(torch.nn.Module):
        mask_pad: torch.Tensor,
    ) -> torch.Tensor:
        for layer in self.encoders:
-            xs, chunk_masks, _, _ = ckpt.checkpoint(
-                layer.__call__, xs, chunk_masks, pos_emb, mask_pad
+            xs, chunk_masks, _, _ = torch.utils.checkpoint.checkpoint(
+                layer.__call__, xs, chunk_masks, pos_emb, mask_pad, use_reentrant=False
            )
        return xs