fixing gradient calculation