Hello, thank you for implemantation. I would like to train it on my custom data? What should i change to make it work?