def generate_code(self) -> Callable:
        """
        Obtain the correct function for the given input and output
        definitions.

        :return: The callable to be used as the transformation.
        """

        def convert_apply_convert_reshape(inp, _):
            inp_as_tensor = torch.from_numpy(inp)
            # N, H, W, C -> N, C, H, W
            inp_as_tensor = inp_as_tensor.permute([0, 3, 1, 2])
            res = self.module(inp_as_tensor)

            # N, C, H, W -> N, H, W, C
            res_as_np: np.ndarray = res.numpy()
            return res_as_np.transpose((0, 2, 3, 1))

        def convert_apply_reshape(inp, _):
            inp_as_tensor = torch.from_numpy(inp)
            # N, H, W, C -> N, C, H, W
            inp_as_tensor = inp_as_tensor.permute([0, 3, 1, 2])

            res = self.module(inp_as_tensor)
            return res

        def apply_convert_reshape(inp, _):
            res = self.module(inp)

            # N, C, H, W -> N, H, W, C
            res_as_np: np.ndarray = res.numpy()
            return res_as_np.transpose((0, 2, 3, 1))

        def convert_apply_convert(inp, _):
            inp_as_tensor = torch.from_numpy(inp)
            res = self.module(inp_as_tensor)
            return res.numpy()

        def convert_apply(inp, _):
            inp_as_tensor = torch.from_numpy(inp)
            res = self.module(inp_as_tensor)
            return res

        def apply_convert(inp, _):
            res = self.module(inp)
            return res.numpy()

        def apply(inp, _):
            device = inp.device
            return self.module(inp).to(device, non_blocking=True)

        # (input_type, output_type, smart_reshape) -> func
        func_table = {
            ("numpy", "numpy", True): convert_apply_convert_reshape,
            ("numpy", "torch", True): convert_apply_reshape,
            ("torch", "numpy", True): apply_convert_reshape,
            ("numpy", "numpy", False): convert_apply_convert,
            ("numpy", "torch", False): convert_apply,
            ("torch", "numpy", False): apply_convert,
            ("torch", "torch", True): apply,
            ("torch", "torch", False): apply,
        }

        return func_table[(self.input_type, self.output_type, self.smart_reshape)]

        Obtain the correct function for the given input and output
        definitions.

        :return: The callable to be used as the transformation.
        
avalanche/benchmarks/utils/ffcv_support/ffcv_transform_utils.py