def __init__(
        self,
        module: torch.nn.Module,
        expected_out_type: Union[
            np.dtype, torch.dtype, Literal["as_previous"]
        ] = "as_previous",
        expected_shape: Union[Tuple[int, ...], Literal["as_previous"]] = "as_previous",
        smart_reshape: bool = True,
    ):
        """
        Creates an instance of a SmartModuleWrapper.

        :param module: The module to use. The module must be able to process
            the inputs in batches.
        :param expected_out_type: The expected type of the output. Default to `as_previous`.
        :param expected_shape: The expected shape of the output. Default to `as_previous`.
        :param smart_reshape: If True, will try to compute the proper shape conversion
            when the input is NumPy and the shape suggests that an image is being passed.
        """
        super().__init__()
        self.module = module
        self.expected_out_type = expected_out_type
        self.expected_shape = expected_shape
        self.input_type = "numpy"
        self.output_type = "numpy"
        self.smart_reshape = smart_reshape

        Creates an instance of a SmartModuleWrapper.

        :param module: The module to use. The module must be able to process
            the inputs in batches.
        :param expected_out_type: The expected type of the output. Default to `as_previous`.
        :param expected_shape: The expected shape of the output. Default to `as_previous`.
        :param smart_reshape: If True, will try to compute the proper shape conversion
            when the input is NumPy and the shape suggests that an image is being passed.
        
avalanche/benchmarks/utils/ffcv_support/ffcv_transform_utils.py