Подсчет градиентов в градиентном спуске

Помогите найти ошибку в формуле градиентов. Вот, как их считаю я:

dw2 = - (2 / n) * np.dot(self.A[1].T,  2 * (self.A[2] - Y))
db2 = - (2 / n) * np.sum(self.A[2] - Y)
delta = self.A[1] - self.A[2]
dw1 = - (2 / n) * np.dot(X.T, delta)
db1 = - (2 / n) * np.sum(self.A[1] - self.A[2])
self.params['w2'] -= self.lr * dw2
self.params['b2'] -= self.lr * db2
self.params['w1'] -= self.lr * dw1
self.params['b1'] -= self.lr * db1

A1 - выход скрытого слоя нейросети, A2 - ответ после сигмоиды (A - это словарь с ключами 1 и 2). Использую lr = 0.01 (пробовал от 0.1 до 1e-5), архитектура сети: linear(out_size = 8) -> ReLU ->linear(out_size = 1) -> sigmoid. При обучении график потерь выглядит как-то так: