[WIP] Gradients working
Fix tests for gradients