Skip to content

Commit f190d99

Browse files
committed
Update if nvjpeg not found
1 parent 62e072a commit f190d99

File tree

2 files changed

+10
-8
lines changed

2 files changed

+10
-8
lines changed

benchmarks/encoding.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,6 @@
11
import os
22
import platform
33
import statistics
4-
import tarfile
5-
import tempfile
6-
import urllib.request
74

85
import torch
96
import torch.utils.benchmark as benchmark
@@ -52,7 +49,7 @@ def run_benchmark(batch):
5249
stmt=stmt,
5350
setup="import torchvision",
5451
globals={"batch_input": batch_input},
55-
label=f"Image Encoding",
52+
label="Image Encoding",
5653
sub_label=f"{device.upper()} ({strat}): {stmt}",
5754
description=f"{size} images",
5855
num_threads=num_threads,

torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
#include "encode_jpegs_cuda.h"
22
#if !NVJPEG_FOUND
3+
namespace vision {
4+
namespace image {
35
std::vector<torch::Tensor> encode_jpegs_cuda(
4-
const std::vector<torch::Tensor>& images,
6+
const std::vector<torch::Tensor>& decoded_images,
57
const int64_t quality) {
68
TORCH_CHECK(
79
false, "encode_jpegs_cuda: torchvision not compiled with nvJPEG support");
810
}
11+
} // namespace image
12+
} // namespace vision
913
#else
1014

1115
#include <ATen/ATen.h>
@@ -111,9 +115,10 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
111115
CUDAJpegEncoder::CUDAJpegEncoder(const torch::Device& target_device)
112116
: original_device{torch::kCUDA, torch::cuda::current_device()},
113117
target_device{target_device},
114-
stream{at::cuda::getStreamFromPool(
115-
false,
116-
target_device.has_index() ? target_device.index() : 0)} {
118+
stream{
119+
target_device.has_index()
120+
? at::cuda::getStreamFromPool(false, target_device.index())
121+
: at::cuda::getStreamFromPool(false)} {
117122
nvjpegStatus_t status;
118123
status = nvjpegCreateSimple(&nvjpeg_handle);
119124
TORCH_CHECK(

0 commit comments

Comments
 (0)