diff --git a/.github/workflows/train.yaml b/.github/workflows/train.yaml
index bc1f8001..b57d7aee 100644
--- a/.github/workflows/train.yaml
+++ b/.github/workflows/train.yaml
@@ -3,20 +3,34 @@ on: [push]
 jobs:
   run:
     runs-on: [ubuntu-latest]
-    container: docker://dvcorg/cml-py3:latest
+    container: docker://dvcorg/cml:latest
     steps:
       - uses: actions/checkout@v2
+
+      - name: Set safe directory
+        run: git config --global --add safe.directory /__w/githubactionsml/githubactionsml
+
+      - name: Set up Git
+        run: |
+          git config --global user.email "github-actions[bot]@users.noreply.github.com"
+          git config --global user.name "github-actions[bot]"
+
       - name: cml_run
         env:
           repo_token: ${{ secrets.GITHUB_TOKEN }}
         run: |
           pip install -r requirements.txt
           dvc repro 
-
           git fetch --prune
-          dvc metrics diff --show-md master > report.md
-
+          echo "## DVC metrics comparison" >> report.md
+          dvc metrics diff main >> report.md
           # Add figure to the report
-          echo "## Validating results by region"
+          echo "## Validating results by region" >> report.md
           cml-publish by_region.png --md >> report.md
-          cml-send-comment report.md
\ No newline at end of file
+          cml-send-comment report.md
+
+      - name: Commit DVC changes
+        run: |
+          git add dvc.lock
+          git commit -m "Update DVC lock file" || echo "No changes to commit"
+          git push
diff --git a/requirements.txt b/requirements.txt
index 51f3da2a..f97c9c99 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 wget
-sklearn
+scikit-learn
 pandas
 seaborn
 matplotlib
diff --git a/train.py b/train.py
index e37f0eb5..5b5ea1ed 100644
--- a/train.py
+++ b/train.py
@@ -48,6 +48,6 @@
 # Bar plot by region
 
 sns.set_color_codes("dark")
-ax = sns.barplot(x="region", y="pred_accuracy", data=df, palette = "Greens_d")
+ax = sns.barplot(x="region", y="pred_accuracy", hue="region", data=df, palette="Greens_d", legend=False)
 ax.set(xlabel="Region", ylabel = "Model accuracy")
 plt.savefig("by_region.png",dpi=80)