Call nvjpegJpegStateDestroy when bailing out

pytorch · Jan 23, 2021 · e485656 · e485656
1 parent a07f53a
commit e485656
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchvision/csrc/io/image/cuda/readjpeg_cuda.cpp b/torchvision/csrc/io/image/cuda/readjpeg_cuda.cpp
@@ -92,6 +92,7 @@ torch::Tensor decodeJPEG_cuda(const torch::Tensor& data, ImageReadMode mode) {
         outputFormat = NVJPEG_OUTPUT_RGB;
         outputComponents = 3;
       } else {
+        nvjpegJpegStateDestroy(nvjpeg_state);
         TORCH_CHECK(
             false, "The provided mode is not supported for JPEG files on GPU");
       }
@@ -108,6 +109,7 @@ torch::Tensor decodeJPEG_cuda(const torch::Tensor& data, ImageReadMode mode) {
       break;
     default:
       // CMYK as input might work with nvjpegDecodeParamsSetAllowCMYK()
+      nvjpegJpegStateDestroy(nvjpeg_state);
       TORCH_CHECK(
           false, "The provided mode is not supported for JPEG files on GPU");
   }