Skip to content

Commit

Permalink
Fix optional chaining for batch size calculation in PreTrainedModel
Browse files Browse the repository at this point in the history
  • Loading branch information
chenyuankun committed Dec 2, 2024
1 parent 9584263 commit 4f08a3d
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/models.js
Original file line number Diff line number Diff line change
Expand Up @@ -1757,7 +1757,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) {
Expand Down

0 comments on commit 4f08a3d

Please sign in to comment.