Fix gradient calculation; Avoid unnecessary allocations