debug

sbintuitions · Aug 9, 2024 · abe4f88 · abe4f88
1 parent 39f98a3
commit abe4f88
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/jmteb/evaluators/reranking/evaluator.py b/src/jmteb/evaluators/reranking/evaluator.py
@@ -151,7 +151,7 @@ def _compute_metrics(
 
         with tqdm.tqdm(total=len(query_dataset), desc="Reranking docs") as pbar:
             if torch.cuda.is_available():
-                if dist.is_available():
+                if dist.is_available() and dist.is_torchelastic_launched():
                     device = f"cuda:{dist.get_rank()}"
                 else:
                     device = "cuda"

diff --git a/src/jmteb/evaluators/retrieval/evaluator.py b/src/jmteb/evaluators/retrieval/evaluator.py
@@ -160,7 +160,7 @@ def _compute_metrics(
                 doc_embeddings_chunk = doc_embeddings[offset : offset + self.doc_chunk_size]
 
                 if torch.cuda.is_available():
-                    if dist.is_available():
+                    if dist.is_available() and dist.is_torchelastic_launched():
                         device = f"cuda:{dist.get_rank()}"
                     else:
                         device = "cuda"