Fix tensorflow import in train.py

[vexing-shusher]

Hello! Thank you for your work. Could you please remove the "import tensorflow as tf" line from train.py? It does not seem to be used anywhere in the code.