Always convert features to correct dtype

lightly-ai · Dec 15, 2023 · c0d09ee · c0d09ee
1 parent 453f9bf
commit c0d09ee
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/lightly/utils/benchmarking/knn_classifier.py b/lightly/utils/benchmarking/knn_classifier.py
@@ -105,7 +105,8 @@ def training_step(self, batch, batch_idx) -> None:
         images, targets = batch[0], batch[1]
         features = self.model.forward(images).flatten(start_dim=1)
         if self.normalize:
-            features = F.normalize(features, dim=1).to(self.feature_dtype)
+            features = F.normalize(features, dim=1)
+        features = features.to(self.feature_dtype)
         self._train_features.append(features.cpu())
         self._train_targets.append(targets.cpu())
 
@@ -116,7 +117,8 @@ def validation_step(self, batch, batch_idx) -> None:
         images, targets = batch[0], batch[1]
         features = self.model.forward(images).flatten(start_dim=1)
         if self.normalize:
-            features = F.normalize(features, dim=1).to(self.feature_dtype)
+            features = F.normalize(features, dim=1)
+        features = features.to(self.feature_dtype)
         predicted_classes = knn_predict(
             feature=features,
             feature_bank=self._train_features_tensor,