def threeD_dataset (X, y, time_steps = 1):
    Xs, ys = [], []
    
    for i in range(len(X)-time_steps):
        v = X[i:i+time_steps, :]
        Xs.append(v)
        ys.append(y[i+time_steps])
        
    return np.array(Xs), np.array(ys)


TIME_STEPS = 50

X_test, y_test = threeD_dataset(test_x_norm, y_test, TIME_STEPS)
X_train, y_train = threeD_dataset(train_x_norm, y_train, TIME_STEPS)