relu simpler backward pass
MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -6,11 +6,12 @@ from tinygrad.tensor import Function
class ReLU(Function):
def forward(self, x):
- self.save_for_backward(x)
- return x.unary_op(UnaryOps.RELU)
+ ret = x.unary_op(UnaryOps.RELU)
+ self.save_for_backward(ret)
+ return ret
def backward(self, grad_output):
- return self.saved_tensors[0].unary_op(UnaryOps.SIGN).unary_op(UnaryOps.RELU).binary_op(BinaryOps.MUL, grad_output)
+ return self.saved_tensors[0].unary_op(UnaryOps.SIGN).binary_op(BinaryOps.MUL, grad_output)
class Log(Function):
def forward(self, x):

relu simpler backward pass