github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

hlb_cifar10 torch version

Browse Source
This commit is contained in:
George Hotz
2023-02-11 18:04:40 -08:00
parent 9057d98d36
commit 191c76cfd7
3 changed files with 103 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tinygrad/nn/__init__.py
View File

@@ -1,5 +1,6 @@
from tinygrad.tensor import Tensor
# TODO: BatchNorm2D -> 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"