iris = load_iris() data = pd.DataFrame(iris.data, columns=iris.feature_names) y = iris.target y = np.identity(3, dtype=np.int64)[y] X_train, X_test, y_train, y_test = train_test_split(data.values, y, train_size=0.67, shuffle=True) print(f"X_train: {X_train.shape}, X_test: {X_test.shape}, y_train: {y_train.shape}, y_test: {y_test.shape}") X_train = torch.from_numpy(X_train).float() y_train = torch.