Forgot the return statement
Merged softmax
../../../models/mlp.py
Working on parallelism
More attempt at parallelism