diff --git a/.github/workflows/cml.yaml b/.github/workflows/cml.yaml
index fa5e65d..a0e16c8 100644
--- a/.github/workflows/cml.yaml
+++ b/.github/workflows/cml.yaml
@@ -7,16 +7,22 @@ jobs:
       - uses: actions/checkout@v3
         with:
           ref: ${{ github.event.pull_request.head.sha }}
-      - uses: iterative/setup-cml@v1
+      - uses: iterative/setup-cml@v2
       - uses: actions/setup-python@v2
         with:
           python-version: '3.x'
+      - name: Install requirements
+        run: |
+          pip3 install -r requirements.txt
+      - name: Get data
+        run: |
+          python3 get_data.py
       - name: Train model
+        run: |
+          python3 train.py
+      - name: Create PR with metrics
         env:
           repo_token: ${{ secrets.GITHUB_TOKEN }}
         run: |
-          pip3 install -r requirements.txt
-          python train.py
-
-          # Create CML report
-          cml comment create --pr --publish metrics.txt
+          cml pr create .
+          cml comment create --target=pr --publish metrics.txt
diff --git a/requirements.txt b/requirements.txt
index ad78221..5dce312 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,2 @@
-setuptools
-scikit-learn
-matplotlib
+scikit-learn==1.3.2
+matplotlib==3.8.2
diff --git a/train.py b/train.py
index c08584f..e3346c7 100644
--- a/train.py
+++ b/train.py
@@ -2,7 +2,7 @@
 import os
 
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.metrics import plot_confusion_matrix
+from sklearn.metrics import ConfusionMatrixDisplay
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -29,5 +29,5 @@
     outfile.write(metrics)
 
 # Plot it
-disp = plot_confusion_matrix(clf, X_test, y_test, normalize="true", cmap=plt.cm.Blues)
+disp = ConfusionMatrixDisplay.from_estimator(clf, X_test, y_test, normalize="true", cmap=plt.cm.Blues)
 plt.savefig("plot.png")