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

tensor inference (#6156)

Browse Source 
* tensor inference

* test is even better name
This commit is contained in:
George Hotz
2024-08-18 00:19:28 -07:00
committed by GitHub
parent 9db2d0d5c6
commit 17a043edad
3 changed files with 12 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
test/test_tensor.py
View File

@@ -642,11 +642,11 @@ class TestTrainMode(unittest.TestCase):
assert not Tensor.training
class TestInferenceMode(unittest.TestCase):
def test_inference_mode(self):
def test_inference(self):
x = Tensor(x_init, requires_grad=True)
m = Tensor(m_init, requires_grad=True)
W = Tensor(W_init, requires_grad=True)
with Tensor.inference_mode():
with Tensor.test():
tmp = x.mul(m)
mm = tmp.matmul(W)
out = mm.relu()
@@ -663,7 +663,7 @@ class TestInferenceMode(unittest.TestCase):
x = Tensor(x_init, requires_grad=True)
m = Tensor(m_init, requires_grad=True)
W = Tensor(W_init, requires_grad=True)
@Tensor.inference_mode()
@Tensor.test()
def f(x, m, W):
tmp = x.mul(m)
mm = tmp.matmul(W)