Thank you for your source code. Can you help me about loss function and how to train model in pytorch ?