Unverified Commit 4d10ad56 authored by Baber Abbasi's avatar Baber Abbasi Committed by GitHub
Browse files

batch_schedular bug in Collator (#1229)

* auto-batch requires len of iter

* handle case when batch_size="auto:N"
parent a12ef445
......@@ -96,7 +96,8 @@ class VLLM(LM):
trust_remote_code=trust_remote_code,
tokenizer_revision=tokenizer_revision,
)
self.batch_size = batch_size
self.batch_size = "auto" if batch_size.startswith("auto:") else batch_size
self._max_gen_toks = max_gen_toks
@property
......
......@@ -920,6 +920,7 @@ class Collator:
```
"""
arr = []
_iter = tuple(_iter)
for i, x in enumerate(_iter):
arr.append(x)
if len(arr) == (fn(i, _iter) if fn else n):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment