diff --git a/src/mozmlops/templates/template_flow.py b/src/mozmlops/templates/template_flow.py
index 7d95063..5df70b4 100644
--- a/src/mozmlops/templates/template_flow.py
+++ b/src/mozmlops/templates/template_flow.py
@@ -108,7 +108,8 @@ def train(self):
 
         # Load the Iris dataset
         iris = load_iris()
-        X, y = iris.data, iris.target
+        print('is the error here')
+        X, y = iris['data'], iris['target']
 
         # Split the dataset into training and testing sets
         X_train, X_test, y_train, y_test = train_test_split(