I have added a new code to train.py

[giliolimirko]

Description

New code has been added

How has this been tested?

Please explain how you have tested the new changes.

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes (make test-trigger and make test-all-components)
• I have successfully run the E2E tests
• I have updated any relevant documentation to reflect my changes

[giliolimirko]

/gcbrun