fix for #114 -- incorrect shape used in preload (#115)

ml-explore · Aug 30, 2024 · 9266a62 · 9266a62
1 parent ab94ffc
commit 9266a62
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/Libraries/LLM/Evaluate.swift b/Libraries/LLM/Evaluate.swift
@@ -179,7 +179,7 @@ public struct TokenIterator: Sequence, IteratorProtocol {
         // prepare the prompt in chunks if larger than the prefill size
         while y.size > parameters.prefillStepSize {
             _ = model(
-                y[..<parameters.prefillStepSize, .newAxis], cache: cache.isEmpty ? nil : cache)
+                y[.newAxis, ..<parameters.prefillStepSize], cache: cache.isEmpty ? nil : cache)
             eval(cache)
             y = y[parameters.prefillStepSize...]
         }