diff --git a/torchvision/csrc/io/image/cpu/encode_jpeg.cpp b/torchvision/csrc/io/image/cpu/encode_jpeg.cpp index c84ad37005d..1ba730f0998 100644 --- a/torchvision/csrc/io/image/cpu/encode_jpeg.cpp +++ b/torchvision/csrc/io/image/cpu/encode_jpeg.cpp @@ -18,12 +18,12 @@ using namespace detail; torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { // Define compression structures and error handling - struct jpeg_compress_struct cinfo; - struct torch_jpeg_error_mgr jerr; + struct jpeg_compress_struct cinfo {}; + struct torch_jpeg_error_mgr jerr {}; // Define buffer to write JPEG information to and its size unsigned long jpegSize = 0; - uint8_t* jpegBuf = NULL; + uint8_t* jpegBuf = nullptr; cinfo.err = jpeg_std_error(&jerr.pub); jerr.pub.error_exit = torch_jpeg_error_exit; @@ -34,7 +34,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { * We need to clean up the JPEG object and the buffer. */ jpeg_destroy_compress(&cinfo); - if (jpegBuf != NULL) { + if (jpegBuf != nullptr) { free(jpegBuf); } @@ -92,16 +92,10 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) { jpeg_destroy_compress(&cinfo); torch::TensorOptions options = torch::TensorOptions{torch::kU8}; - auto outTensor = torch::empty({(long)jpegSize}, options); - - // Copy memory from jpeg buffer, since torch cannot get ownership of it via - // `from_blob` - auto outPtr = outTensor.data_ptr(); - std::memcpy(outPtr, jpegBuf, sizeof(uint8_t) * outTensor.numel()); - - free(jpegBuf); - - return outTensor; + auto out_tensor = + torch::from_blob(jpegBuf, {(long)jpegSize}, ::free, options); + jpegBuf = nullptr; + return out_tensor; } #endif