diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultCodec.java b/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultCodec.java
index 24ea6c5f8c..66b9e15554 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultCodec.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultCodec.java
@@ -361,6 +361,11 @@ public final class DefaultCodec implements Codec {
       if (outputBufferIndex == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED) {
         outputFormat =
             convertToFormat(mediaCodec.getOutputFormat(), isDecoder, configurationFormat.metadata);
+        // The raw audio decoder incorrectly sets the channel count for output format to stereo.
+        if (isDecoder && Objects.equals(configurationFormat.sampleMimeType, MimeTypes.AUDIO_RAW)) {
+          outputFormat =
+              outputFormat.buildUpon().setChannelCount(configurationFormat.channelCount).build();
+        }
       }
       return false;
     }
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java b/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
index ee8f67e18b..0a16123fe2 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/ExoAssetLoaderBaseRenderer.java
@@ -219,6 +219,7 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
       }
       inputFormat = overrideInputFormat(checkNotNull(formatHolder.format));
       onInputFormatRead(inputFormat);
+      // TODO: b/332708880 - Bypass MediaCodec for raw audio input.
       shouldInitDecoder =
           assetLoaderListener.onTrackAdded(
               inputFormat, SUPPORTED_OUTPUT_TYPE_DECODED | SUPPORTED_OUTPUT_TYPE_ENCODED);
@@ -262,14 +263,6 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
         if (decoderOutputFormat == null) {
           return false;
         }
-
-        // TODO: b/332708880 - The raw decoder caps number of audio channels to stereo. Remove this
-        //  workaround and bypass MediaCodec for raw audio instead.
-        if (checkNotNull(inputFormat.sampleMimeType).equals(MimeTypes.AUDIO_RAW)) {
-          decoderOutputFormat =
-              decoderOutputFormat.buildUpon().setChannelCount(inputFormat.channelCount).build();
-        }
-
         outputFormat = overrideOutputFormat(decoderOutputFormat);
       } else {
         // TODO(b/278259383): Move surface creation out of video sampleConsumer. Init decoder and