train_stage2 meet problem: index is out of bounds for dimension

[bingoyesbingoyes]

Steps:  16%|█▌        | 3732/24000 [7:56:34<39:55:30,  7.09s/it, lr=1e-5, step_loss=0.0711, td=4.81s]
Steps:  16%|█▌        | 3733/24000 [7:56:42<41:35:07,  7.39s/it, lr=1e-5, step_loss=0.0711, td=4.81s]
Steps:  16%|█▌        | 3733/24000 [7:56:42<41:35:07,  7.39s/it, lr=1e-5, step_loss=0.0306, td=6.20s]09/26/2024 03:57:51 - ERROR - root - Failed to execute the training process: index 301 is out of bounds for dimension 0 with size 301

Steps:  16%|█▌        | 3733/24000 [7:56:42<43:08:09,  7.66s/it, lr=1e-5, step_loss=0.0306, td=6.20s]
Moviepy - Done !
Moviepy - video ready ./exp_output/stage2/validation/9000_1_1.mp4
Checkpoint saved at ./exp_output/stage2/modules/net-9000.pth
Checkpoint saved at ./exp_output/stage2/modules/net-9500.pth

When I'm training on stage2, the index is out bounds of dimension.