Patched the train.py script to accept tensors as input

6 jobs for updates in 13 minutes and 11 seconds (queued for 2 seconds)