fixing issues with cuda #2

sensein · Sep 19, 2024 · e3d5acc · e3d5acc
1 parent ba1ffac
commit e3d5acc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/.github/workflows/e2c-runner-tests-310.yaml b/.github/workflows/e2c-runner-tests-310.yaml
@@ -125,7 +125,7 @@ jobs:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
       run: >
-        cd senselab && poetry run pytest -n auto \
+        cd senselab && poetry run pytest \
           --rootdir=$POETRY_CACHE_DIR/pytest \
           --basetemp=$POETRY_CACHE_DIR/pytest/temp \
           --junitxml=pytest.xml \

diff --git a/src/senselab/text/tasks/embeddings_extraction/huggingface.py b/src/senselab/text/tasks/embeddings_extraction/huggingface.py
@@ -78,6 +78,9 @@ def extract_text_embeddings(
         device, _ = _select_device_and_dtype(
             user_preference=device, compatible_devices=[DeviceType.CUDA, DeviceType.CPU]
         )
+
+        print(f"Using device: {device}")
+
         # Load tokenizer and model
         tokenizer = cls._get_tokenizer(model=model)
         ssl_model = cls._load_model(model=model, device=device)
@@ -87,13 +90,15 @@ def extract_text_embeddings(
         # Process each piece of text individually
         for text in pieces_of_text:
             # Tokenize sentence
-            encoded_input = tokenizer(text, return_tensors="pt").to(device)
+            encoded_input = tokenizer(text, return_tensors="pt").to(device.value)
 
             # Compute token embeddings
             with torch.no_grad():
                 model_output = ssl_model(**encoded_input, output_hidden_states=True)
                 hidden_states = model_output.hidden_states
-                concatenated_hidden_states = torch.cat([state.unsqueeze(0) for state in hidden_states], dim=0)
+                concatenated_hidden_states = torch.cat(
+                    [state.to(device.value).unsqueeze(0) for state in hidden_states], dim=0
+                )
                 embeddings.append(concatenated_hidden_states.squeeze())
 
         return embeddings