replace with tensor op (#3099)

2026-09-01 07:46:07 +00:00 · 2024-01-12 14:13:40 -05:00
parent f3a50b4e40
commit 0c49d38ba7
1 changed files with 1 additions and 1 deletions
@@ -51,7 +51,7 @@ class Transformer:
    maxlen_eye = Tensor.eye(x.shape[1])
    maxlen_eye = maxlen_eye.unsqueeze(0).expand([bs, *maxlen_eye.shape])

-    onehot_feat = x[:, :, None] == Tensor.arange(self.syms).reshape((1,1) + (self.syms,))
+    onehot_feat = x.one_hot(self.syms)

    onehot = maxlen_eye.cat(onehot_feat, dim=2).flatten(end_dim=1)