Reapply "Fix memory leak in decode_webp (pytorch#8712)" (pytorch#8723)

This reverts commit 7d077f1.
NicolasHug · Nov 12, 2024 · 6282f2c · 6282f2c
1 parent 7d077f1
commit 6282f2c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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});
 }