github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-02-18 02:21:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix batchnorm not realizing

Browse Source
This commit is contained in:
George Hotz
2023-02-27 10:19:54 -08:00
parent c9252d38b2
commit 2f17d151b3
4 changed files with 24 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/nn/optim.py
View File

@@ -10,6 +10,8 @@ class Optimizer:
x.requires_grad = True
self.params : List[Tensor] = [x for x in params if x.requires_grad]
self.buffers : List[Tensor] = [x for x in params if not x.requires_grad] # buffers are still realized
self.realize()
# TODO: this probably shouldn't change the gradients, just the ones used by the optimizer
def clipnorm(self, amount=1):
@@ -24,7 +26,7 @@ class Optimizer:
def realize(self, extra=None):
# TODO: corealize
for p in extra + self.params if extra is not None else self.params:
for p in extra + self.params + self.buffers if extra is not None else self.params + self.buffers:
p.realize()
class SGD(Optimizer):