diff --git a/.travis.yml b/.travis.yml
index 712fc84..b97d816 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -5,10 +5,11 @@ python:
 services:
     - docker
 
-install:
-    - pip install flake8
-    - docker build . -t baselines-test
+env:
+  - DOCKER_SUFFIX=py36-nompi
+  - DOCKER_SUFFIX=py36-mpi
 
+install: "" # so travis doesn't do pip install requirements.txt
 script:
-    - flake8 . --show-source --statistics
-    - docker run baselines-test pytest -v .
+  - flake8 . --show-source --statistics
+  - docker build -f test.dockerfile.${DOCKER_SUFFIX} -t baselines-test .
diff --git a/Dockerfile b/test.dockerfile.py36-mpi
similarity index 66%
rename from Dockerfile
rename to test.dockerfile.py36-mpi
index 49a9c79..d56966b 100644
--- a/Dockerfile
+++ b/test.dockerfile.py36-mpi
@@ -1,6 +1,6 @@
 FROM python:3.6
 
-# RUN apt-get -y update && apt-get -y install git wget python-dev python3-dev libopenmpi-dev python-pip zlib1g-dev cmake python-opencv
+RUN apt-get -y update && apt-get -y install ffmpeg libopenmpi-dev
 ENV CODE_DIR /root/code
 
 COPY . $CODE_DIR/baselines
diff --git a/test.dockerfile.py36-nompi b/test.dockerfile.py36-nompi
new file mode 100644
index 0000000..636d97c
--- /dev/null
+++ b/test.dockerfile.py36-nompi
@@ -0,0 +1,16 @@
+FROM python:3.6
+
+RUN apt-get -y update && apt-get -y install ffmpeg
+ENV CODE_DIR /root/code
+
+COPY . $CODE_DIR/baselines
+WORKDIR $CODE_DIR/baselines
+
+# Clean up pycache and pyc files
+RUN rm -rf __pycache__ && \
+    find . -name "*.pyc" -delete && \
+    pip install tensorflow && \
+    pip install -e .[test]
+
+
+CMD /bin/bash