From 52f0081e779e54b4f98adf6a2f677f68cd57acc2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 5 Nov 2025 12:49:01 -0500 Subject: [PATCH] use where instead of mul in Embedding (#13112) --- tinygrad/nn/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index c8884146d3..5d5ced5c32 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -323,7 +323,7 @@ class Embedding: if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}") big_shp = idx.shape+(self.vocab_sz, self.embed_sz) arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), self.weight.expand(big_shp) - return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype) + return (arange == idx).where(vals, 0).sum(-2, dtype=vals.dtype) class LSTMCell: """