(items, dim=None)
| 39 | |
| 40 | |
| 41 | def maybe_stack(items, dim=None): |
| 42 | to_stack = [item for item in items if item is not None] |
| 43 | if not to_stack: |
| 44 | return None |
| 45 | elif len(to_stack) == 1: |
| 46 | return to_stack[0].unsqueeze(dim) |
| 47 | else: |
| 48 | return torch.stack(to_stack, dim) |
| 49 | |
| 50 | |
| 51 | def accumulate_logprobs(d, keys_and_logprobs): |