diff --git a/crossfit/op/base.py b/crossfit/op/base.py
index 5a564f69..14e213f3 100644
--- a/crossfit/op/base.py
+++ b/crossfit/op/base.py
@@ -80,7 +80,7 @@ def __call__(self, data, *args, partition_info=None, **kwargs):
             self.teardown()
             return output
 
-        # self.setup()
+        self.setup()
 
         if self.pre is not None:
             params = inspect.signature(self.pre).parameters
diff --git a/docker/ci/Dockerfile b/docker/ci/Dockerfile
index 2557a75e..b92cc2e1 100644
--- a/docker/ci/Dockerfile
+++ b/docker/ci/Dockerfile
@@ -1,9 +1,8 @@
-FROM nvcr.io/nvidia/pytorch:23.10-py3
+FROM nvcr.io/nvidia/nemo:23.10
 
 COPY . /tmp/crossfit/
 RUN cd /tmp/crossfit && \
     python3 -m pip install .[pytorch-dev] && \
-    python3 -m pip install -r requirements/gpu.txt --no-cache-dir --extra-index-url https://pypi.nvidia.com && \
     python3 -m pip install beir && \
     rm -r /tmp/crossfit
 
diff --git a/requirements/gpu.txt b/requirements/gpu.txt
deleted file mode 100644
index 1685a93f..00000000
--- a/requirements/gpu.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-cudf-cu12>=23.12
-dask-cuda>=23.12
-dask-cudf-cu12>=23.12