(batch, pad_id, seq_length)
| 115 | |
| 116 | |
| 117 | def pad_batch(batch, pad_id, seq_length): |
| 118 | context_lengths = [] |
| 119 | for tokens in batch: |
| 120 | context_length = len(tokens) |
| 121 | if context_length < seq_length: |
| 122 | tokens.extend([pad_id] * (seq_length - context_length)) |
| 123 | context_lengths.append(context_length) |
| 124 | return batch, context_lengths |
| 125 | |
| 126 | |
| 127 | def forward_step( |