refactor gradient
fix for compatability with new AutoEncoder