drop cudf 23.12 and use nemo contianer

rapidsai · Jan 15, 2024 · 324a64d · 324a64d
1 parent 130ac37
commit 324a64d
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 6 deletions.
diff --git 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
@@ -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