zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-09-01 04:14:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Feat: update pre-commit rev (#432)

Browse Source
This commit is contained in:
Nripesh Niketan
2024-02-11 19:23:27 +04:00
committed by GitHub
parent f45a1ab83c
commit f1ef378a58
3 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
llms/mlx_lm/tuner/trainer.py
View File

@@ -80,9 +80,9 @@ def iterate_batches(dataset, tokenizer, batch_size, max_seq_length, train=False)
for j in range(batch_size):
truncated_length = min(lengths[j], max_seq_length)
batch_arr[j, :truncated_length] = batch[j][:truncated_length]
lengths[
j
] = truncated_length # Update lengths to match truncated lengths
lengths[j] = (
truncated_length # Update lengths to match truncated lengths
)
batch = mx.array(batch_arr)
yield batch[:, :-1], batch[:, 1:], mx.array(lengths)