From 9cd41ac746b581b8c7022ac63da760f17e3cea28 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 29 Apr 2026 16:45:06 -0700 Subject: [PATCH 1/2] import umap when needed --- viscy/representation/engine.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/viscy/representation/engine.py b/viscy/representation/engine.py index 6c2d243d9..6f679d009 100644 --- a/viscy/representation/engine.py +++ b/viscy/representation/engine.py @@ -7,7 +7,6 @@ from lightning.pytorch import LightningModule from pytorch_metric_learning.losses import NTXentLoss from torch import Tensor, nn -from umap import UMAP from viscy.data.typing import TrackingIndex, TripletSample from viscy.representation.contrastive import ContrastiveEncoder @@ -183,6 +182,8 @@ def _log_step_samples(self, batch_idx, samples, stage: Literal["train", "val"]): output_list.extend(detach_sample(samples, self.log_samples_per_batch)) def log_embedding_umap(self, embeddings: Tensor, tag: str): + from umap import UMAP + _logger.debug(f"Computing UMAP for {tag} embeddings.") umap = UMAP(n_components=2) embeddings_np = embeddings.detach().cpu().numpy() From 625e92b341cc2f29a850f0e20ac6dc4862a0f7f7 Mon Sep 17 00:00:00 2001 From: Ivan Ivanov Date: Wed, 29 Apr 2026 16:54:14 -0700 Subject: [PATCH 2/2] format --- viscy/representation/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/viscy/representation/engine.py b/viscy/representation/engine.py index 6f679d009..00dea509a 100644 --- a/viscy/representation/engine.py +++ b/viscy/representation/engine.py @@ -183,7 +183,7 @@ def _log_step_samples(self, batch_idx, samples, stage: Literal["train", "val"]): def log_embedding_umap(self, embeddings: Tensor, tag: str): from umap import UMAP - + _logger.debug(f"Computing UMAP for {tag} embeddings.") umap = UMAP(n_components=2) embeddings_np = embeddings.detach().cpu().numpy()