torch_tensor_from_array_int16_3d Subroutine

  subroutine torch_tensor_from_array_int16_3d(tensor, data_in, layout, &
                                                        c_device_type, device_index, requires_grad)
    use, intrinsic :: iso_c_binding, only : c_bool, c_float, c_int, c_int64_t, c_loc
    use, intrinsic :: iso_fortran_env, only : int16

    ! output tensor
    type(torch_tensor), intent(out) :: tensor !! Returned tensor

    ! inputs
    integer(kind=int16), intent(in), target :: data_in(:,:,:)   !! Input data that tensor will point at
    integer, intent(in)        :: layout(3) !! Control order of indices
    integer(c_int), intent(in) :: c_device_type    !! Device type the tensor will live on (`torch_kCPU` or `torch_kCUDA`)
    integer(c_int), optional, intent(in) :: device_index    !! device index to use for `torch_kCUDA` case
    logical, optional, intent(in) :: requires_grad  !! Whether gradients need to be computed for the created tensor

    ! local data
    integer(c_int64_t)        :: c_tensor_shape(3)           !! Shape of the tensor
    integer(c_int), parameter :: c_dtype = torch_kInt16 !! Data type
    integer(c_int64_t)        :: strides(3)                  !! Strides for accessing data
    integer(c_int), parameter :: ndims = 3                   !! Number of dimension of input data
    integer                   :: i
    integer(c_int)            :: device_index_value
    logical :: requires_grad_value  !! Whether gradients need to be computed for the created tensor

    ! Process optional arguments
    if (present(device_index)) then
      device_index_value = device_index
    else if (c_device_type == torch_kCPU) then
      device_index_value = -1
    else
      device_index_value = 0
    endif

    if (.not. present(requires_grad)) then
      requires_grad_value = .false.
    else
      requires_grad_value = requires_grad
    end if

    c_tensor_shape = shape(data_in)

    strides(layout(1)) = 1
    do i = 2, ndims
      strides(layout(i)) = strides(layout(i - 1)) * c_tensor_shape(layout(i - 1))
    end do

    tensor%p = torch_from_blob_c(c_loc(data_in), ndims, c_tensor_shape,        &
                                 strides, c_dtype, c_device_type,              &
                                 device_index_value,                           &
                                 logical(requires_grad_value, c_bool))

  end subroutine torch_tensor_from_array_int16_3d