def forward(self, *args, **kwargs) -> Union[torch.Tensor, Tuple[torch.Tensor]]:
        r"""Forward method for the `SplitInferenceModule`.

        This method processes the input by splitting specified keyword arguments along a given dimension, running the
        underlying module on each split, and then concatenating the results. The splitting is controlled by the
        `split_size` and `split_dim` parameters specified during initialization.

        Args:
            *args (`Any`):
                Positional arguments that are passed directly to the `module` without modification.
            **kwargs (`Dict[str, torch.Tensor]`):
                Keyword arguments passed to the underlying `module`. Only keyword arguments whose names match the
                entries in `input_kwargs_to_split` and are of type `torch.Tensor` will be split. The remaining keyword
                arguments are passed unchanged.

        Returns:
            `Union[torch.Tensor, Tuple[torch.Tensor]]`:
                The outputs obtained from `SplitInferenceModule` are the same as if the underlying module was inferred
                without it.
                - If the underlying module returns a single tensor, the result will be a single concatenated tensor
                along the same `split_dim` after processing all splits.
                - If the underlying module returns a tuple of tensors, each element of the tuple will be concatenated
                along the `split_dim` across all splits, and the final result will be a tuple of concatenated tensors.
        """
        split_inputs = {}

        # 1. Split inputs that were specified during initialization and also present in passed kwargs
        for key in list(kwargs.keys()):
            if key not in self.input_kwargs_to_split or not torch.is_tensor(kwargs[key]):
                continue
            split_inputs[key] = torch.split(kwargs[key], self.split_size, self.split_dim)
            kwargs.pop(key)

        # 2. Invoke forward pass across each split
        results = []
        for split_input in zip(*split_inputs.values()):
            inputs = dict(zip(split_inputs.keys(), split_input))
            inputs.update(kwargs)

            intermediate_tensor_or_tensor_tuple = self.module(*args, **inputs)
            results.append(intermediate_tensor_or_tensor_tuple)

        # 3. Concatenate split restuls to obtain final outputs
        if isinstance(results[0], torch.Tensor):
            return torch.cat(results, dim=self.split_dim)
        elif isinstance(results[0], tuple):
            return tuple([torch.cat(x, dim=self.split_dim) for x in zip(*results)])
        else:
            raise ValueError(
                "In order to use the SplitInferenceModule, it is necessary for the underlying `module` to either return a torch.Tensor or a tuple of torch.Tensor's."
            )
Forward method for the `SplitInferenceModule`.

        This method processes the input by splitting specified keyword arguments along a given dimension, running the
        underlying module on each split, and then concatenating the results. The splitting is controlled by the
        `split_size` and `split_dim` parameters specified during initialization.

        Args:
            *args (`Any`):
                Positional arguments that are passed directly to the `module` without modification.
            **kwargs (`Dict[str, torch.Tensor]`):
                Keyword arguments passed to the underlying `module`. Only keyword arguments whose names match the
                entries in `input_kwargs_to_split` and are of type `torch.Tensor` will be split. The remaining keyword
                arguments are passed unchanged.

        Returns:
            `Union[torch.Tensor, Tuple[torch.Tensor]]`:
                The outputs obtained from `SplitInferenceModule` are the same as if the underlying module was inferred
                without it.
                - If the underlying module returns a single tensor, the result will be a single concatenated tensor
                along the same `split_dim` after processing all splits.
                - If the underlying module returns a tuple of tensors, each element of the tuple will be concatenated
                along the `split_dim` across all splits, and the final result will be a tuple of concatenated tensors.
        
src/diffusers/pipelines/free_noise_utils.py