From 6282f2c2d96f8b8e5578aa584d99ba26a74cd4bc Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 12 Nov 2024 13:15:46 +0000 Subject: [PATCH] Reapply "Fix memory leak in decode_webp (#8712)" (#8723) This reverts commit 7d077f131217dc03813d97d7524ea3aeba7dd7e1. --- torchvision/csrc/io/image/cpu/decode_webp.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchvision/csrc/io/image/cpu/decode_webp.cpp b/torchvision/csrc/io/image/cpu/decode_webp.cpp index b202473c039..0a9ff9ddbce 100644 --- a/torchvision/csrc/io/image/cpu/decode_webp.cpp +++ b/torchvision/csrc/io/image/cpu/decode_webp.cpp @@ -44,10 +44,12 @@ torch::Tensor decode_webp( auto decoded_data = decoding_func(encoded_data_p, encoded_data_size, &width, &height); + TORCH_CHECK(decoded_data != nullptr, "WebPDecodeRGB[A] failed."); + auto deleter = [decoded_data](void*) { WebPFree(decoded_data); }; auto out = torch::from_blob( - decoded_data, {height, width, num_channels}, torch::kUInt8); + decoded_data, {height, width, num_channels}, deleter, torch::kUInt8); return out.permute({2, 0, 1}); }