geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

argfix axis

Browse Source
This commit is contained in:
David Hou
2024-02-21 23:34:06 -08:00
parent 5e8b6f8d10
commit 6cac3c99e2
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+1 -1
View File
@@ -928,7 +928,7 @@ class Tensor:
shape = tuple(s if ax in argfix(axis) else 1 for ax, s in enumerate(self.shape))
x = self - mean.reshape(shape)
if weight: x = x * weight.reshape(shape)
ret = x.mul(invstd.reshape(shape) if len(invstd.shape) == len(axis) else invstd)
ret = x.mul(invstd.reshape(shape) if len(invstd.shape) == len(argfix(axis)) else invstd)
return (ret + bias.reshape(shape)) if bias else ret
def dropout(self, p=0.5) -> Tensor: