megatron.schedules.forward_step

megatron.schedules.forward_step(forward_step_func, data_iterator, model, input_tensor, forward_data_store, timers, collect_non_loss_data=False)

Forward step for passed-in model.

If first stage, input tensor is obtained from data_iterator, otherwise passed-in input_tensor is used.

Returns output tensor.