Enable gradient checkpointing

acestep/models/lyrics_utils/lyric_encoder.py

@@ -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