set_training_mode Module Subroutine

  module subroutine set_training_mode(this, mode_store, layer_indices)
    !! Put the network in training mode.
    implicit none

    ! Arguments
    class(network_type), intent(inout) :: this
    !! Instance of network
    logical, dimension(:), allocatable, intent(out), optional :: mode_store
    !! Optional array to store the training mode of each layer
    integer, dimension(:), intent(in), optional :: layer_indices
    !! Optional array of layer indices to set to training mode.
    !! If not provided, all layers will be set to training mode.

    ! Local variables
    integer :: l
    !! Loop index

    if(.not.allocated(this%model)) return
    if(present(mode_store)) allocate(mode_store(this%num_layers))
    do l = 1, this%num_layers
       if(present(mode_store)) mode_store(l) = this%model(l)%layer%inference
       this%model(l)%layer%inference = .false.
       if(present(layer_indices))then
          if(any(layer_indices.eq.l))then
             this%model(l)%layer%inference = .false.
          end if
       else
          this%model(l)%layer%inference = .false.
       end if
    end do

  end subroutine set_training_mode