Fixes

flash/image/classification/integrations/baal/data.py

@@ -141,7 +141,10 @@ def train_dataloader(self) -> "DataLoader":
         if self.has_labelled_data and self.val_split:
             self.val_dataloader = self._val_dataloader
 
-        return self.labelled.train_dataloader()
+        if self.has_labelled_data:
+            return self.labelled.train_dataloader()
+        # Return a dummy dataloader, will be replaced by the loop
+        return DataLoader(["dummy"])
 
     def _val_dataloader(self) -> "DataLoader":
         self.labelled._val_input = train_val_split(self._dataset, self.val_split)[1]

flash/image/classification/integrations/baal/loop.py

@@ -185,6 +185,7 @@ def _reset_dataloader_for_stage(self, running_state: RunningStage):
             if is_overridden(dataloader_name, self.trainer.datamodule)
             else None
         )
+
         if dataloader:
             if _PL_GREATER_EQUAL_1_5_0:
                 setattr(

flash/image/face_detection/model.py

@@ -71,7 +71,7 @@ def __init__(
             learning_rate=learning_rate,
             optimizer=optimizer,
             lr_scheduler=lr_scheduler,
-            output_transform=FaceDetectionOutputTransform,
+            output_transform=FaceDetectionOutputTransform(),
         )
 
     @staticmethod

flash/image/face_detection/output_transform.py

@@ -37,6 +37,5 @@ def per_batch_transform(batch: Any) -> Any:
         # preds: list of torch.Tensor(N, 5) as x1, y1, x2, y2, score
         preds = [preds[preds[:, 5] == batch_idx, :5] for batch_idx in range(len(preds))]
         preds = ff.utils.preprocess.adjust_results(preds, scales, paddings)
-        batch[DataKeys.PREDS] = preds
 
-        return batch
+        return preds