def METHOD_NAME(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
    result = self(batch[DataKeys.INPUT])
    result[DataKeys.INPUT] = default_uncollate(batch[DataKeys.INPUT])
    return default_uncollate(result)