Added profiling code. Need to profile and speed up training.