set_gradients Module Subroutine

  module subroutine set_gradients(this, gradients)
    !! Set gradients
    implicit none

    ! Arguments
    class(network_type), intent(inout) :: this
    !! Instance of network
    real(real32), dimension(..), intent(in) :: gradients
    !! Gradients

    ! Local variables
    integer :: l, start_idx, end_idx
    !! Loop index

    start_idx = 0
    end_idx   = 0
    do l = 1, this%num_layers
       select type(current => this%model(l)%layer)
       class is(learnable_layer_type)
          start_idx = end_idx + 1
          end_idx = end_idx + current%num_params
          select rank(gradients)
          rank(0)
             call current%set_gradients(gradients)
          rank(1)
             call current%set_gradients(gradients(start_idx:end_idx))
          end select
       end select
    end do

  end subroutine set_gradients