You can even strip out the autodiff and write an explicit backwards pass, training time went down from 40s to 5s for me.