train

Hello, author, first of all, thank you for your work. I encountered this problem when I tried to reproduce your code. Can you give me some advice?