From 939655abb226eb43bbc75e8fcb4f13223ccd1335 Mon Sep 17 00:00:00 2001
From: Nicolas Hug <contact@nicolas-hug.com>
Date: Mon, 17 Mar 2025 17:48:32 +0000
Subject: [PATCH] Simplify UniqueAVFrame handling in audio decoder

---
 src/torchcodec/decoders/_core/VideoDecoder.cpp | 18 +++++++-----------
 src/torchcodec/decoders/_core/VideoDecoder.h   |  2 +-
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/src/torchcodec/decoders/_core/VideoDecoder.cpp b/src/torchcodec/decoders/_core/VideoDecoder.cpp
index 9871db64..e68bb356 100644
--- a/src/torchcodec/decoders/_core/VideoDecoder.cpp
+++ b/src/torchcodec/decoders/_core/VideoDecoder.cpp
@@ -1361,14 +1361,8 @@ void VideoDecoder::convertAudioAVFrameToFrameOutputOnCPU(
       static_cast<AVSampleFormat>(avFrameStream.avFrame->format);
   AVSampleFormat desiredSampleFormat = AV_SAMPLE_FMT_FLTP;
 
-  UniqueAVFrame convertedAVFrame;
-  if (sourceSampleFormat != desiredSampleFormat) {
-    convertedAVFrame = convertAudioAVFrameSampleFormat(
-        avFrameStream.avFrame, sourceSampleFormat, desiredSampleFormat);
-  }
-  const UniqueAVFrame& avFrame = (sourceSampleFormat != desiredSampleFormat)
-      ? convertedAVFrame
-      : avFrameStream.avFrame;
+  const UniqueAVFrame avFrame = convertAudioAVFrameSampleFormat(
+      avFrameStream.avFrame, sourceSampleFormat, desiredSampleFormat);
 
   AVSampleFormat format = static_cast<AVSampleFormat>(avFrame->format);
   TORCH_CHECK(
@@ -1395,11 +1389,13 @@ void VideoDecoder::convertAudioAVFrameToFrameOutputOnCPU(
 }
 
 UniqueAVFrame VideoDecoder::convertAudioAVFrameSampleFormat(
-    const UniqueAVFrame& avFrame,
+    UniqueAVFrame& avFrame,
     AVSampleFormat sourceSampleFormat,
-    AVSampleFormat desiredSampleFormat
+    AVSampleFormat desiredSampleFormat) {
+  if (sourceSampleFormat == desiredSampleFormat) {
+    return std::move(avFrame);
+  }
 
-) {
   auto& streamInfo = streamInfos_[activeStreamIndex_];
   const auto& streamMetadata =
       containerMetadata_.allStreamMetadata[activeStreamIndex_];
diff --git a/src/torchcodec/decoders/_core/VideoDecoder.h b/src/torchcodec/decoders/_core/VideoDecoder.h
index f72f31ab..7711df2a 100644
--- a/src/torchcodec/decoders/_core/VideoDecoder.h
+++ b/src/torchcodec/decoders/_core/VideoDecoder.h
@@ -404,7 +404,7 @@ class VideoDecoder {
       torch::Tensor& outputTensor);
 
   UniqueAVFrame convertAudioAVFrameSampleFormat(
-      const UniqueAVFrame& avFrame,
+      UniqueAVFrame& avFrame,
       AVSampleFormat sourceSampleFormat,
       AVSampleFormat desiredSampleFormat);