github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

with Tensor.train() (#1935)

Browse Source 
* add with.train

* remove the rest TODOs

* fix pyflake

* fix pyflake error

* fix mypy
This commit is contained in:
Yixiang Gao
2023-09-28 20:02:31 -05:00
committed by GitHub
parent 10f0dc0c85
commit 094d3d71be
14 changed files with 305 additions and 317 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
test/test_schedule.py
View File

@@ -154,12 +154,11 @@ class TestSchedule(unittest.TestCase):
#@unittest.skip("may want to reconsider this")
def test_fold_batchnorm(self):
Tensor.training = True
img = Tensor.empty(1,32,4,4)
bn = nn.BatchNorm2d(32, track_running_stats=False)
out = bn(img)
check_schedule(out, 3)
Tensor.training = False
with Tensor.train():
img = Tensor.empty(1,32,4,4)
bn = nn.BatchNorm2d(32, track_running_stats=False)
out = bn(img)
check_schedule(out, 3)
def test_fold_conv_relu(self):
c1 = nn.Conv2d(3,16,3)