Forward step.
(data_iterator, model)
| 198 | |
| 199 | |
| 200 | def forward_step(data_iterator, model): |
| 201 | """Forward step.""" |
| 202 | args = get_args() |
| 203 | timers = get_timers() |
| 204 | |
| 205 | # Get the batch. |
| 206 | timers("batch-generator").start() |
| 207 | tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator) |
| 208 | timers("batch-generator").stop() |
| 209 | |
| 210 | output_tensor = model(tokens, position_ids, attention_mask, labels=labels) |
| 211 | |
| 212 | return output_tensor, partial(loss_func, loss_mask) |
| 213 | |
| 214 | |
| 215 | def valid_forward_step(data_iterator, model): |
nothing calls this directly
no test coverage detected