class SGD:
    def __init__(self, params, lr=0.01):
        self.params = params
        self.lr = lr

    def step(self):
        for p in self.params:
            if p.grad is not None:
                # Если градиент раздулся (например, из-за батча), 
                # мы приводим его к форме весов (p.data)
                g = p.grad
                if g.shape != p.data.shape:
                    # Суммируем по осям, чтобы форма совпала
                    g = g.sum(axis=0, keepdims=True)
                
                p.data -= self.lr * g

    def zero_grad(self):
        for p in self.params:
            p.grad = None