Fix a few other loading.

liuliu · Mar 15, 2024 · e9b31c4 · e9b31c4
1 parent 2e85afa
commit e9b31c4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nnc/Store.swift b/nnc/Store.swift
@@ -3401,7 +3401,8 @@ private let q8pDecodeJitWithExternalEager:
     }
     guard
       TensorShape(dims: params.dim).reduce(1, *) == numberOfElements
-        && (numberOfElements % 4) == 0  // We support non-block size length for q8p only.
+        && (numberOfElements % (256 * 4)) == 0
+        && (blockSize % (256 * 4)) == 0 // We support non-block size length for q8p only.
     else {
       let mappedData = store.loadBytes(offset: offset, length: length)
       defer {
@@ -3773,7 +3774,8 @@ private let q8pDecodeJitWithExternalOnDemand:
     }
     guard
       TensorShape(dims: params.dim).reduce(1, *) == numberOfElements
-        && (numberOfElements % 4) == 0  // We support non-block size length for q8p only.
+        && (numberOfElements % (256 * 4)) == 0
+        && (blockSize % (256 * 4)) == 0 // We support non-block size length for q8p only.
     else {
       let mappedData = store.loadBytes(offset: offset, length: length)
       defer {