sample['target'] from a dataloader batch returns the wrong z-shape

[edyoshikun]

We were expecting that after setting an HCSDataModule().test_dataloader()

for i, sample in enumerate(HCSDataModule().test_dataloader()):
     break
sample['target'].shape # returns( torch.Size([64, 2, 5, 512, 512]))

The expected behavior should be that sample['target'] is 1.

Note:
When predicting, I see that we select the middle of the stack, but the default setting for the target should be always be one

def on_before_batch_transfer(self, batch: Sample, dataloader_idx: int) -> Sample:
        predicting = False
        if self.trainer:
            if self.trainer.predicting:
                predicting = True
        if predicting or isinstance(batch, torch.Tensor):
            # skipping example input array
            return batch
        if self.target_2d:
            # slice the center during training or testing
            z_index = self.z_window_size // 2
            batch["target"] = batch["target"][:, :, slice(z_index, z_index + 1)]
        return batch

[ziw-liu]

@edyoshikun fixed by #51?