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

Question

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?