From bad4489a39bd8d34287eba0ea61324adf0e4d692 Mon Sep 17 00:00:00 2001
From: Omkar Kabde <omkarkabde@gmail.com>
Date: Sun, 15 Feb 2026 04:19:27 +0530
Subject: [PATCH] refactor output tracing in `univnet`

---
 src/transformers/models/univnet/modeling_univnet.py | 12 +++---------
 1 file changed, 3 insertions(+), 9 deletions(-)

diff --git a/src/transformers/models/univnet/modeling_univnet.py b/src/transformers/models/univnet/modeling_univnet.py
index a7436eb7bd1e..c72f83f1629c 100644
--- a/src/transformers/models/univnet/modeling_univnet.py
+++ b/src/transformers/models/univnet/modeling_univnet.py
@@ -20,7 +20,7 @@
 
 from ...modeling_outputs import ModelOutput
 from ...modeling_utils import PreTrainedModel
-from ...utils import auto_docstring, logging
+from ...utils import auto_docstring, can_return_tuple, logging
 from .configuration_univnet import UnivNetConfig
 
 
@@ -467,6 +467,7 @@ def __init__(self, config: UnivNetConfig):
         # Initialize weights and apply final processing
         self.post_init()
 
+    @can_return_tuple
     @auto_docstring
     def forward(
         self,
@@ -474,9 +475,8 @@ def forward(
         noise_sequence: torch.FloatTensor | None = None,
         padding_mask: torch.FloatTensor | None = None,
         generator: torch.Generator | None = None,
-        return_dict: bool | None = None,
         **kwargs,
-    ) -> tuple[torch.FloatTensor] | UnivNetModelOutput:
+    ) -> UnivNetModelOutput:
         r"""
         noise_sequence (`torch.FloatTensor`, *optional*):
             Tensor containing a noise sequence of standard Gaussian noise. Can be batched and of shape `(batch_size,
@@ -516,8 +516,6 @@ def forward(
          [1, 140288]
          ```
         """
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
         # Resolve batch sizes for noise_sequence and spectrogram
         spectrogram_batched = input_features.dim() == 3
         if not spectrogram_batched:
@@ -582,10 +580,6 @@ def forward(
             # Padding is always contiguous and added on the right
             waveform_lengths = torch.sum(padding_mask, dim=1)
 
-        if not return_dict:
-            outputs = (waveform, waveform_lengths)
-            return outputs
-
         return UnivNetModelOutput(
             waveforms=waveform,
             waveform_lengths=waveform_lengths,