[IMP] Some changes to adding gradients