Location:
datasets.py
train_test_split()
line 206
Current: dimension = dimension if not None else X.shape[1]
Correct: dimension = dimension if dimension is not None else X.shape[1]
Since None is always False, the current behavior will leave dimension as None, rather than changing it to X.shape[1] when dimension=None as I assume was intended.