derivative of softmax is indepedent of max

chenyuxyz · Oct 12, 2024 · f7b8c41 · f7b8c41
1 parent f79e05c
commit f7b8c41
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
@@ -1682,7 +1682,7 @@ def std_mean(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False,
 
   def _softmax(self, axis, dtype:Optional[DTypeLike]=None):
     x = self.cast(dtype) if dtype is not None else self
-    m = x - x.max(axis=axis, keepdim=True)
+    m = x - x.max(axis=axis, keepdim=True).detach()
     e = m.exp()
     return m, e, e.sum(axis=axis, keepdim=True)