From ba49e357a1deb9bf9b04a425eca5b3227e7b42ba Mon Sep 17 00:00:00 2001
From: Vibhu Jawa <vibhujawa@gmail.com>
Date: Thu, 3 Oct 2024 17:53:40 -0600
Subject: [PATCH 1/2] fix conda environment (#93)

Signed-off-by: Vibhu Jawa <vibhujawa@gmail.com>
---
 conda/environments/cuda_dev.yaml | 4 +++-
 requirements/cuda12x.txt         | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/conda/environments/cuda_dev.yaml b/conda/environments/cuda_dev.yaml
index e092385..f82d0f0 100644
--- a/conda/environments/cuda_dev.yaml
+++ b/conda/environments/cuda_dev.yaml
@@ -20,13 +20,15 @@ dependencies:
   - pytest
   - pytest-benchmark
   - pytest-cov>=2
-  - pytorch>=2.0
+  - pytorch>=2.0=*cuda*
+  - tiktoken
   - transformers
   - tokenizers
   - torchvision
   - scikit-learn>=0.23.1
   - scipy
   - setuptools>=61.0.0
+  - sentencepiece
   - sentence-transformers
   - pip:
     - pytrec-eval
diff --git a/requirements/cuda12x.txt b/requirements/cuda12x.txt
index 54e1c2d..24632b9 100644
--- a/requirements/cuda12x.txt
+++ b/requirements/cuda12x.txt
@@ -11,3 +11,4 @@ curated-transformers>=1.0
 bitsandbytes>=0.30
 sentence-transformers>=2.0
 sentencepiece
+tiktoken

From d7e2643ce2520452cfd7475620b63e65034c2dc9 Mon Sep 17 00:00:00 2001
From: Vibhu Jawa <vibhujawa@gmail.com>
Date: Tue, 8 Oct 2024 22:22:14 -0600
Subject: [PATCH 2/2] Fix logging of GPU ID  (#95)

* Fix logging

* Fix torch CI GPU issue

* Fix CI GPU install problem

* One more attempt to fix  torch CI GPU issue

* One more chance Fix torch CI GPU issue

* Fix typo
---
 ci/test_gpu.sh      | 15 ++++++---------
 crossfit/op/base.py |  5 ++++-
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/ci/test_gpu.sh b/ci/test_gpu.sh
index 606f19c..6523110 100755
--- a/ci/test_gpu.sh
+++ b/ci/test_gpu.sh
@@ -18,18 +18,15 @@ fi
 echo "Installing pytorch,transformers and pytest to the environment for crossfit tests..."
 mamba install \
   cuda-version=$CUDA_VERSION \
-  conda-forge::pytorch \
-  conda-forge::transformers \
-  conda-forge::pytest \
+  "pytorch>=2.*=*cuda*" \
+  transformers \
+  pytest \
+  sentence-transformers  \
+  sentencepiece \
   -c conda-forge \
-  --override-channels \
+  -c nvidia \
   --yes
 
-# Have to install sentence-transformers from pip
-# because conda-forge leads to a torch vision conflict
-# which leads to it being installed on CPUs
-pip3 install sentence-transformers sentencepiece
-
 # Install the crossfit package in editable mode with test dependencies
 pip3 install -e '.[test]'
 # Running tests
diff --git a/crossfit/op/base.py b/crossfit/op/base.py
index d9ea28d..2e9557a 100644
--- a/crossfit/op/base.py
+++ b/crossfit/op/base.py
@@ -26,7 +26,10 @@ def __init__(self, pre=None, cols=False, keep_cols=None):
         self.pre = pre
         self.cols = cols
         self.keep_cols = keep_cols or []
-        self.worker_name = getattr(self.get_worker(), "name", 0)
+
+    @property
+    def worker_name(self):
+        return getattr(self.get_worker(), "name", 0)
 
     def setup(self):
         pass