Fixing gradient bug (#13)