Re Guarin

lightly-ai · Nov 20, 2023 · 5ee6086 · 5ee6086
1 parent 4c442a2
commit 5ee6086
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/lightly/utils/benchmarking/metric_callback.py b/lightly/utils/benchmarking/metric_callback.py
@@ -60,4 +60,7 @@ def _append_metrics(
         self, metrics_dict: Dict[str, List[float]], trainer: Trainer
     ) -> None:
         for name, value in trainer.callback_metrics.items():
-            metrics_dict.setdefault(name, []).append(float(value))
+            if isinstance(value, float) or (  # type: ignore # We can't rely on PyTorchLightning's type annotations.
+                isinstance(value, Tensor) and value.numel() == 1
+            ):
+                metrics_dict.setdefault(name, []).append(float(value))
diff --git a/lightly/utils/embeddings_2d.py b/lightly/utils/embeddings_2d.py
@@ -66,8 +66,8 @@ def transform(self, X: NDArray[np.float32]) -> NDArray[np.float32]:
             raise ValueError("PCA not fitted yet. Call fit() before transform().")
         X = X.astype(np.float32)
         X = X - self.mean + self.eps
-        transformed = X.dot(self.w)[:, : self.n_components]
-        return np.asarray(transformed, dtype=np.float32)
+        transformed: NDArray[np.float32] = X.dot(self.w)[:, : self.n_components]
+        return np.asarray(transformed)
 
 
 def fit_pca(