-1 is not a valid index

frost-beta · May 22, 2024 · 342b092 · 342b092
1 parent b23bfdb
commit 342b092
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lib/nn/layers/transformer.ts b/lib/nn/layers/transformer.ts
@@ -87,7 +87,7 @@ export class MultiHeadAttention extends Module {
     values = values.reshape(B, S, numHeads, -1).transpose(0, 2, 1, 3);
 
     // Dimensions are [batch x numHeads x sequence x hiddenDim].
-    const scale = mx.array(Math.sqrt(1 / queries.shape[-1]), queries.dtype);
+    const scale = Math.sqrt(1 / queries.shape[queries.shape.length - 1]);
     let scores = mx.matmul(mx.multiply(queries, scale), keys);
     if (mask)
       scores = mx.add(scores, mask.astype(scores.dtype));
@@ -417,4 +417,4 @@ export class Transformer extends Module {
     const memory = this.encoder.forward(src, srcMask);
     return this.decoder.forward(tgt, memory, tgtMask, memoryMask);
   }
-}
+}