gitea-mirror/F5-TTS
1
0
Fork 0
mirror of https://github.com/SWivid/F5-TTS.git synced 2026-07-28 13:26:59 -07:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #859 from ZhikangNiu/main

Browse Source 
fix #858 and pass use_reentrant explicitly in checkpoint_activation mode
This commit is contained in:
Yushen CHEN
2025-03-15 16:23:50 +08:00
committed by GitHub
parent 297755fac3 f017815083
commit 2bd39dd813
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/f5_tts/model/backbones/dit.py
+3 -1
View File
@@ -219,7 +219,9 @@ class DiT(nn.Module):
for block in self.transformer_blocks:
if self.checkpoint_activations:
x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope)
# if you have question, please check: https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint
# After PyTorch 2.4, we must pass the use_reentrant explicitly
x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False)
else:
x = block(x, t, mask=mask, rope=rope)