[Projeto] Erro ao criar o modelo ML

Olá, estou tendo um pouco de dificuldade pra acompanhar o curso devido não estar acostumado ao VSCode. Acabei tendo um erro ao tentar criar o modelo de machine learning e não consegui identificá-lo para corrigir até o momento. Agradeço a quem puder me ajudar. A seguir está meu código do treina_modelo.py e logo abaixo dele a resposta do terminal. Acredito que não seja nada complexo, mas não consegui identificar. Desde já agradeço.

Código do treina_modelo.py:

from ucimlrepo import fetch_ucirepo

heart_disease = fetch_ucirepo(id=45)
dados = heart_disease.data.features
dados['doenca'] = (heart_disease.data.targets > 0) * 1

X = dados.drop(columns = 'doenca')
y = dados['doenca']

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 432, stratify = y)

import xgboost as xgb
modelo = xgb.XGBClassifier(objective = 'binary:logistic')

modelo.fit(X_train, y_train)

preds = modelo.predict(X_test)

from sklearn.metrics import accuracy_score
acuracia = accuracy_score(y_test, preds)
print(f'A acurácia do modelo é {acuracia:.2f}')

Resposta no terminal:

Traceback (most recent call last):
  File "c:\Users\jopim\Downloads\ALURA_CURSOS\treina_modelo.py", line 11, in <module>
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 432, stratify = y)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\jopim\Downloads\ALURA_CURSOS\.venv\Lib\site-packages\sklearn\utils\_param_validation.py", line 213, in wrapper      
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\jopim\Downloads\ALURA_CURSOS\.venv\Lib\site-packages\sklearn\model_selection\_split.py", line 2782, in train_test_split
    arrays = indexable(*arrays)
             ^^^^^^^^^^^^^^^^^^
  File "C:\Users\jopim\Downloads\ALURA_CURSOS\.venv\Lib\site-packages\sklearn\utils\validation.py", line 514, in indexable
    check_consistent_length(*result)
  File "C:\Users\jopim\Downloads\ALURA_CURSOS\.venv\Lib\site-packages\sklearn\utils\validation.py", line 457, in check_consistent_length
    raise ValueError(
ValueError: Found input variables with inconsistent numbers of samples: [301, 303]