Update if nvjpeg not found

pytorch · Jun 5, 2024 · f190d99 · f190d99
1 parent 62e072a
commit f190d99
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/benchmarks/encoding.py b/benchmarks/encoding.py
@@ -1,9 +1,6 @@
 import os
 import platform
 import statistics
-import tarfile
-import tempfile
-import urllib.request
 
 import torch
 import torch.utils.benchmark as benchmark
@@ -52,7 +49,7 @@ def run_benchmark(batch):
                         stmt=stmt,
                         setup="import torchvision",
                         globals={"batch_input": batch_input},
-                        label=f"Image Encoding",
+                        label="Image Encoding",
                         sub_label=f"{device.upper()} ({strat}): {stmt}",
                         description=f"{size} images",
                         num_threads=num_threads,

diff --git a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp
@@ -1,11 +1,15 @@
 #include "encode_jpegs_cuda.h"
 #if !NVJPEG_FOUND
+namespace vision {
+namespace image {
 std::vector<torch::Tensor> encode_jpegs_cuda(
-    const std::vector<torch::Tensor>& images,
+    const std::vector<torch::Tensor>& decoded_images,
     const int64_t quality) {
   TORCH_CHECK(
       false, "encode_jpegs_cuda: torchvision not compiled with nvJPEG support");
 }
+} // namespace image
+} // namespace vision
 #else
 
 #include <ATen/ATen.h>
@@ -111,9 +115,10 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
 CUDAJpegEncoder::CUDAJpegEncoder(const torch::Device& target_device)
     : original_device{torch::kCUDA, torch::cuda::current_device()},
       target_device{target_device},
-      stream{at::cuda::getStreamFromPool(
-          false,
-          target_device.has_index() ? target_device.index() : 0)} {
+      stream{
+          target_device.has_index()
+              ? at::cuda::getStreamFromPool(false, target_device.index())
+              : at::cuda::getStreamFromPool(false)} {
   nvjpegStatus_t status;
   status = nvjpegCreateSimple(&nvjpeg_handle);
   TORCH_CHECK(