From 9967e9d4b4c6f729f0c56be1aac67c981f31ecc1 Mon Sep 17 00:00:00 2001 From: Xiangru Lian Date: Mon, 2 Aug 2021 09:37:31 -0700 Subject: [PATCH] [DOCS] Fix fused softmax example script naive softmax implementation (#178) --- python/tutorials/02-fused-softmax.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tutorials/02-fused-softmax.py b/python/tutorials/02-fused-softmax.py index aba4a7835..0f96152e2 100644 --- a/python/tutorials/02-fused-softmax.py +++ b/python/tutorials/02-fused-softmax.py @@ -25,7 +25,7 @@ def naive_softmax(x): # read 2MN elements ; write MN elements z = x - x_max[:, None] # read MN elements ; write MN elements - numerator = torch.exp(x) + numerator = torch.exp(z) # read MN elements ; write M elements denominator = numerator.sum(dim=1) # read 2MN elements ; write MN elements