Patched the train.py script to accept tensors as input

parent 9359b23d
Pipeline #12834 passed with stages
in 13 minutes and 11 seconds