From 6c3626479017a35676ddb7c8f168bed3a682143d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Szymon=20O=C5=BC=C3=B3g?=
 <58388001+SzymonOzog@users.noreply.github.com>
Date: Sat, 2 Mar 2024 16:35:44 +0100
Subject: [PATCH] Improve type hints for optimizer (#3583)

* Improve type hints for optimizer

* lint fix
---
 tinygrad/nn/optim.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tinygrad/nn/optim.py b/tinygrad/nn/optim.py
index 97dcafe312..74f3d1297a 100644
--- a/tinygrad/nn/optim.py
+++ b/tinygrad/nn/optim.py
@@ -22,6 +22,8 @@ class Optimizer:
     # NOTE: in extra is too late for most of the params due to issues with assign
     Tensor.corealize(extra + self.params + self.buffers if extra is not None else self.params + self.buffers)
 
+  def step(self) -> None: raise NotImplementedError
+
 class SGD(Optimizer):
   def __init__(self, params: List[Tensor], lr=0.001, momentum=0, weight_decay=0.0, nesterov=False):
     super().__init__(params, lr)