Fix optional chaining for batch size calculation in PreTrainedModel (#…

…1063) Co-authored-by: chenyuankun <chenyuankun@joyy.com>
huggingface · Dec 2, 2024 · aead8a9 · aead8a9
1 parent 11db949
commit aead8a9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/models.js b/src/models.js
@@ -1814,7 +1814,7 @@ export class PreTrainedModel extends Callable {
             const dtype = session?.config?.kv_cache_dtype ?? 'float32';
             const empty = (dtype === 'float16') ? new Uint16Array() : [];
 
-            const batch_size = (decoderFeeds[this.main_input_name] ?? decoderFeeds.attention_mask).dims?.[0] ?? 1;
+            const batch_size = (decoderFeeds[this.main_input_name] ?? decoderFeeds.attention_mask)?.dims?.[0] ?? 1;
             const shapes = getKeyValueShapes(this.config, { batch_size });
 
             for (const name in shapes) {