Update inference.py

src/itwinai/torch/inference.py

@@ -1,3 +1,4 @@
+import abc
 import os
 from typing import Any, Dict, List, Literal, Optional, Union
 
@@ -103,7 +104,7 @@ def __init__(
             self.config = TrainingConfiguration(**config)
         else:
             self.config = config
-        self.model = self.model
+        self.model = self.model.eval()
         self.strategy = strategy
         self.logger = logger
         self.checkpoints_location = checkpoints_location
@@ -154,7 +155,7 @@ def distribute_model(self) -> None:
             )
         distribute_kwargs = {}
         # Distributed model, optimizer, and scheduler
-        self.model,_,_ = self.strategy.distributed(
+        self.model, _, _ = self.strategy.distributed(
             self.model, None, None, **distribute_kwargs
         )
 
@@ -214,9 +215,10 @@ def execute(
             self.logger.save_hyperparameters(hparams)
 
         all_predictions = dict()
-        for samples_ids, samples in inference_dataset:
+        for ids, (samples_ids, samples) in enumerate(self.inference_dataloader):
             with torch.no_grad():
-                pred = self.model(samples)
+                pred = self.model(samples.to(self.device))
+            pred = self.transform_predictions(pred)
             for idx, pre in zip(samples_ids, pred):
                 # For each item in the batch
                 if pre.numel() == 1:
@@ -263,6 +265,13 @@ def log(
                 batch_idx=batch_idx,
                 **kwargs
             )
+
+    @abc.abstractmethod
+    def transform_predictions(self, batch: Batch) -> Batch:
+        """
+        Post-process the predictions of the torch model (e.g., apply
+        threshold in case of multi-label classifier).
+        """
 
 
 class MulticlassTorchPredictor(TorchPredictor):