train_predict.py

train_predict = model.predict(x_train)
rmse = np.sqrt(mean_squared_error(y_train, train_predict))
r2 = r2_score(y_train, train_predict)

print('Conjunto de treino')
print(rmse)
print(r2)
np.round(train_predict, 0, train_predict)
scatter = plt.scatter(x_train['sepal length (cm)'], x_train['sepal width (cm)'], 
            c=train_predict, cmap=plt.cm.Set1,
            edgecolor='k')
plt.legend(*scatter.legend_elements(),
                    loc="lower left", title="Classes")
plt.show()