1
2X = np.random.randn(1000, 2)
3y = np.random.randint(0, 10, size=1000)
4
5X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.1, stratify=y)
6np.unique(y_train, return_counts=True)
7np.unique(y_val, return_counts=True)
8
9train_dataset = Dataset(X_train, y_train, ...)
10train_loader = DataLoader(train_dataset, ...)
11