fix key conversion

mehta-lab · Oct 17, 2024 · 5e12916 · 5e12916
1 parent c4c0bf9
commit 5e12916
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/viscy/representation/embedding_writer.py b/viscy/representation/embedding_writer.py
@@ -40,7 +40,7 @@ def _move_and_stack_embeddings(
     predictions: Sequence[ContrastivePrediction], key: str
 ) -> NDArray:
     """Move embeddings to CPU and stack them into a numpy array."""
-    return torch.cat([p["features"].cpu() for p in predictions], dim=0).numpy()
+    return torch.cat([p[key].cpu() for p in predictions], dim=0).numpy()
 
 
 class EmbeddingWriter(BasePredictionWriter):