diff --git a/src/torchcodec/_core/Encoder.cpp b/src/torchcodec/_core/Encoder.cpp index 5a6a0d7e..ce33613c 100644 --- a/src/torchcodec/_core/Encoder.cpp +++ b/src/torchcodec/_core/Encoder.cpp @@ -99,8 +99,18 @@ AVSampleFormat findBestOutputSampleFormat(const AVCodec& avCodec) { } // namespace AudioEncoder::~AudioEncoder() { - if (avFormatContext_ && avFormatContext_->pb && !avioContextHolder_) { - avio_close(avFormatContext_->pb); + close_avio(); +} + +void AudioEncoder::close_avio() { + if (avFormatContext_ && avFormatContext_->pb) { + avio_flush(avFormatContext_->pb); + + if (!avioContextHolder_) { + avio_close(avFormatContext_->pb); + // avoids closing again in destructor, which would segfault. + avFormatContext_->pb = nullptr; + } } } @@ -308,6 +318,8 @@ void AudioEncoder::encode() { status == AVSUCCESS, "Error in: av_write_trailer", getFFMPEGErrorStringFromErrorCode(status)); + + close_avio(); } UniqueAVFrame AudioEncoder::maybeConvertAVFrame(const UniqueAVFrame& avFrame) { diff --git a/src/torchcodec/_core/Encoder.h b/src/torchcodec/_core/Encoder.h index 04ed6a13..723849f0 100644 --- a/src/torchcodec/_core/Encoder.h +++ b/src/torchcodec/_core/Encoder.h @@ -33,6 +33,7 @@ class AudioEncoder { void encodeFrame(AutoAVPacket& autoAVPacket, const UniqueAVFrame& avFrame); void maybeFlushSwrBuffers(AutoAVPacket& autoAVPacket); void flushBuffers(); + void close_avio(); UniqueEncodingAVFormatContext avFormatContext_; UniqueAVCodecContext avCodecContext_;