| 579 | }; |
| 580 | |
| 581 | LoopGrad::GradDesc::GradDesc( |
| 582 | Loop* loop, MutableStateSaver* mutable_state_saver, |
| 583 | const VarNodeArray& outgrad_owner) |
| 584 | : m_mutable_state_saver{mutable_state_saver}, |
| 585 | m_fwd_graph_modifier{this, loop->m_desc->sub_graph(), mutable_state_saver}, |
| 586 | m_fwd_counter_provider{loop->m_desc->counter_provider()} { |
| 587 | m_counter_provider_up = CounterProvider::make( |
| 588 | *m_sub_graph, OperatorNodeConfig{loop->comp_node()}.name("counter_up")); |
| 589 | m_counter_var_up = m_counter_provider_up->output(0); |
| 590 | |
| 591 | // add counter var as input to ensure loop executes before grad |
| 592 | add_input(loop->output_counter_var()); |
| 593 | counter_provider()->delta(-1); |
| 594 | |
| 595 | init_virtual_loss(loop->m_desc.get(), outgrad_owner); |
| 596 | init_assignments(); |
| 597 | DescImplBase::set_loop_condition(m_orig_loop_cond_var = (get_counter_var() > 0)); |
| 598 | } |
| 599 | |
| 600 | void LoopGrad::GradDesc::init_virtual_loss( |
| 601 | DescImplBase* fwd_desc, const VarNodeArray& outgrad_owner) { |