scikit-learnでデータ分割するときに配列番号を指定して特定のデータを取ってくる方法がわからない

前提・実現したいこと

scikit-learnでデータ分割するときに、ランダムではなく先頭から順番に分割した後で、
そのデータ分割結果を指定して、検証データ・訓練データとしたいです。

[0]~[4]までそれぞれの分割結果をindexで指定して確認しようとしています。

つまり、順番は固定して分割したデータを、区画ごとに参照したいです。

例

dataset[0] = {'train':[0, 1, 2, 3, 4, 5, 6, 7], 'test':[8, 9]}
dataset[1] = {'train':[2, 3, 4, 5, 6, 7, 8, 9], 'test':[0, 1]}
dataset[2]
dataset[3]
dataset[4]

図の出典元：Qiita記事「比較的少なめのデータで機械学習する時は交差検証 (Cross Validation) をするのです」

発生している問題・エラーメッセージ

現在のコードだとランダムは回避できているのですが、どこに配列番号が格納されているのか、また５回と指定してその候補（例えば、train[2 3 4 5 6 7 8 9] test[0 1]）を確認できない状態です。

該当のソースコード

python
1from sklearn.model_selection import train_test_split
2a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
3a_train, a_test = train_test_split(a,test_size=0.2,random_state=0)
4
5print(a_train)
6print(a_test)
7
8#出力
9[4, 9, 1, 6, 7, 3, 0, 5]
10[2, 8]

試したこと

５分割に関しては以下のコードをもとに分割を試しましたが、この場合、訓練とテスト用データの割合が指定できず、indexを指定するには別のエラーも発生してしまっている状態です。

def cross_validate(session, split_size=5):
  results = []
  kf = KFold(n_splits=split_size)
  for train_idx, val_idx in kf.split(train_x_all, train_y_all):
    train_x = train_x_all[train_idx]
    train_y = train_y_all[train_idx]
    val_x = train_x_all[val_idx]
    val_y = train_y_all[val_idx]
    run_train(session, train_x, train_y)
    results.append(session.run(accuracy, feed_dict={x: val_x, y: val_y}))
  return results

実行コード

from sklearn.model_selection import KFold

def cross_validate(nums, nums_str, split_size=5):
  results = []
  kf = KFold(n_splits=split_size)
  for train_idx, val_idx in kf.split(nums, nums_str):
    print(train_idx, val_idx)
    print(nums[1])
    train_x = nums[train_idx] 
    train_y = nums_str[train_idx]
    val_x = nums[val_idx] 
    val_y = nums_str[val_idx] 

cross_validate([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'], split_size=5)

エラー

[2 3 4 5 6 7 8 9] [0 1]
1
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
in <module>()
     14     val_y = nums_str[val_idx]
     15 
---> 16 cross_validate([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'], split_size=5)


in cross_validate(nums, nums_str, split_size)
      9     print(train_idx, val_idx)
     10     print(nums[1])
---> 11     train_x = nums[train_idx]
     12     train_y = nums_str[train_idx]
     13     val_x = nums[val_idx]

TypeError: only integer scalar arrays can be converted to a scalar index