diff --git a/tests/test_pytorch_env.py b/tests/test_pytorch_env.py
index 5371d526c..23f0aea0f 100644
--- a/tests/test_pytorch_env.py
+++ b/tests/test_pytorch_env.py
@@ -87,7 +87,7 @@ def test_dockerfile(self):
             lines = list(map(lambda x: x.strip(), fin.readlines()))
         dockerfile_text = Path(dockerfile_path).read_text()
         
-        assert lines[1] == "FROM registry.git.vgregion.se/aiplattform/images/pytorch:0.3.2"
+        assert lines[0] == "FROM registry.git.vgregion.se/aiplattform/images/pytorch:0.3.2"
         assert lines[-2] == 'WORKDIR /workspace'
         assert "ADD http://aiav2.vgregion.se/VGC%20Root%20CA%20v2.crt /tmp/vgc_root.der" in lines
         assert "ADD http://aiav2.vgregion.se/VGC%20Issuing%201%20CA%20v2.crt /tmp/vgc_issuing1.der" in lines
diff --git a/tests/test_r_env.py b/tests/test_r_env.py
index 1f3bfb49a..20a4b2b37 100644
--- a/tests/test_r_env.py
+++ b/tests/test_r_env.py
@@ -51,7 +51,7 @@ def test_dockerfile(self):
             lines = list(map(lambda x: x.strip(), fin.readlines()))
         dockerfile_text = Path(dockerfile_path).read_text()
 
-        assert lines[2] == "FROM registry.git.vgregion.se/aiplattform/images/r:0.3.0"
+        assert lines[0] == "FROM registry.git.vgregion.se/aiplattform/images/r:0.3.0"
         assert lines[-5] == 'RUN R -e \'renv::activate()\' && R -e \'renv::restore()\''
         assert lines[-3] == 'USER root'
         assert lines[-2] == 'EXPOSE 8787'
diff --git a/{{ cookiecutter.repo_name }}/Dockerfile b/{{ cookiecutter.repo_name }}/Dockerfile
index dd7c77fed..d49d9ff3b 100644
--- a/{{ cookiecutter.repo_name }}/Dockerfile	
+++ b/{{ cookiecutter.repo_name }}/Dockerfile	
@@ -1,6 +1,6 @@
-{% if cookiecutter.image == 'Tensorflow' %}FROM registry.git.vgregion.se/aiplattform/images/tensorflow:0.2.1{% endif %}
-{% if cookiecutter.image == 'PyTorch' %}FROM registry.git.vgregion.se/aiplattform/images/pytorch:0.3.2{% endif %}
-{% if cookiecutter.image == 'R' %}FROM registry.git.vgregion.se/aiplattform/images/r:0.3.0{% endif %}
+{% if cookiecutter.image == 'Tensorflow' -%}FROM registry.git.vgregion.se/aiplattform/images/tensorflow:0.2.1
+{% elif cookiecutter.image == 'PyTorch' -%}FROM registry.git.vgregion.se/aiplattform/images/pytorch:0.3.2
+{% elif cookiecutter.image == 'R' -%}FROM registry.git.vgregion.se/aiplattform/images/r:0.3.0{%- endif %}
 ENV GID=1000
 ENV UID=1000
 {% if cookiecutter.image == 'R' %}ENV EDITOR_FOCUS_DIR="/workspace"{% endif %}