You first have to initialize a model (as in object oriented programming), and then use it. So use the following instead:
# before the epoch-loop
model = neural_network()
...
# in the loop
output = model(train_X)