how to use the train_x and train_y from sklearn k-fold split generator

I am using the sklearn k-fold generator to split some data 10 times. When I run the code below I expect train_x,train_y,test_x,test_y to contain all 10 splits however only the last split seems to be saved in the variables. I want to either save each split as its own data or have all 10 splits in the variables.

the code is:

def Branin(X,Y):

    PI = 3.14159265359
    a = 1
    b = 5.1 / (4 * pow(PI, 2))
    c = 5 / PI
    r = 6
    s = 10
    t = 1 / (8 * PI)
    
    train_y = a * (X - b * Y ** 2 + c * Y - r) ** 2 + s * (1 - t) * torch.cos(Y) + s
    
    return train_y

x,y=torch.linspace(-5, 10, 40), torch.linspace(0, 15, 40)
X1 = np.column_stack((x, y))
X = torch.from_numpy(X1)
Y = Branin(X[:, 0],X[:, 1])

print(X.shape)
print(Y.shape)

kf = KFold(n_splits=10)
kf.get_n_splits(X)

KFold(n_splits=10, random_state=None, shuffle=False)


for train_index, test_index in kf.split(X,Y):
        train_x, test_x = X[train_index], X[test_index]
        train_y, test_y = Y[train_index], Y[test_index]
        print(test_x)        
#         print("TRAIN:", train_index, "TEST:", test_index)
#example of output for test_x here is the last split in the above test_x#
print("test_x output",test_x)

How can I get test_x for example to be all the splits or say split five once calling it outside the generator?

How many English words
do you know?
Test your English vocabulary size, and measure
how many words do you know
Online Test
Powered by Examplum