diff --git a/Dockerfile b/Dockerfile
index 97a5807..34949b2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -47,7 +47,7 @@ RUN git clone https://github.com/SIPEC-Animal-Data-Analysis/SIPEC.git
 
 WORKDIR /home/user/SIPEC
 
-RUN git checkout main
+RUN git checkout improved_inference
 
 ENV VIRTUAL_ENV=/home/user/SIPEC/env