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

BatchNorm2D -> BatchNorm2d (#558)

Browse Source 
* BatchNorm2D -> BatchNorm2d

* Fix typo
This commit is contained in:
Jacky Lee
2023-02-16 12:31:49 -08:00
committed by GitHub
parent 09315ef34f
commit e172f0087a
9 changed files with 34 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tinygrad/nn/__init__.py
View File

@@ -1,9 +1,8 @@
from tinygrad.tensor import Tensor
# TODO: BatchNorm2D -> BatchNorm2d
class BatchNorm2D:
class BatchNorm2d:
def __init__(self, sz, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
assert affine, "BatchNorm2D is only supported with affine"
assert affine, "BatchNorm2d is only supported with affine"
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz)