Archives/petals
2
0
Fork 0
mirror of https://github.com/bigscience-workshop/petals synced 2024-10-31 09:20:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

check num block kwargs

Browse Source
This commit is contained in:
Your Name 2023-09-06 03:13:08 +03:00
parent 17d278e88a
commit 62e780c054
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/petals/client/sequential_autograd.py
View File

@ -49,11 +49,11 @@ async def sequential_forward(
end_index = end_index if end_index is not None else len(sequence_manager.block_uids)
assert start_index >= 0 and end_index <= len(sequence_manager.block_uids)
assert len(block_kwargs) in (
0,
1,
end_index - start_index,
assert (
len(block_kwargs) in (0, 1, end_index - start_index)
), f"got {end_index - start_index} blocks but {len(block_kwargs)} sets of kwargs"
if len(block_kwargs) == 1:
block_kwargs = block_kwargs * (end_index - start_index)
assert isinstance(inputs, torch.Tensor) and inputs.ndim == 3, f"{type(inputs)}: {inputs.ndim}"
assert is_dummy(prompts) or len(prompts) == len(
sequence_manager.block_uids