Spliting our data into train and test sets and seperating those features that belongs to either label 1 or 0

X0 = [X[i,:] for i in range(len(y)) if y[i] == 0] # returns the row were the corresponding label is 0
X1 = [X[i,:] for i in range(len(y)) if y[i] == 1] # returns the row were the corresponding label is 1

X0_np = np.array(X0) # convert it into an array
X1_np = np.array(X1) 

X0_train = X0_np[:1000,:].T 
X0_test = X0_np[1000:,:].T 

X1_train = X1_np[:1000,:].T  
X1_test = X1_np[1000:,:].T 

X_train = np.hstack([X0_train,X1_train]).T # all training examples
y_train=np.zeros((1,2000))
y_train[0, 1000:] = 1

X_test = np.hstack([X0_test,X1_test]).T # all test examples
y_test=np.zeros((1,1000))
y_test[0, 500:] = 1