File tree 2 files changed +10
-8
lines changed
torchvision/csrc/io/image/cuda
2 files changed +10
-8
lines changed Original file line number Diff line number Diff line change 1
1
import os
2
2
import platform
3
3
import statistics
4
- import tarfile
5
- import tempfile
6
- import urllib .request
7
4
8
5
import torch
9
6
import torch .utils .benchmark as benchmark
@@ -52,7 +49,7 @@ def run_benchmark(batch):
52
49
stmt = stmt ,
53
50
setup = "import torchvision" ,
54
51
globals = {"batch_input" : batch_input },
55
- label = f "Image Encoding" ,
52
+ label = "Image Encoding" ,
56
53
sub_label = f"{ device .upper ()} ({ strat } ): { stmt } " ,
57
54
description = f"{ size } images" ,
58
55
num_threads = num_threads ,
Original file line number Diff line number Diff line change 1
1
#include " encode_jpegs_cuda.h"
2
2
#if !NVJPEG_FOUND
3
+ namespace vision {
4
+ namespace image {
3
5
std::vector<torch::Tensor> encode_jpegs_cuda (
4
- const std::vector<torch::Tensor>& images ,
6
+ const std::vector<torch::Tensor>& decoded_images ,
5
7
const int64_t quality) {
6
8
TORCH_CHECK (
7
9
false , " encode_jpegs_cuda: torchvision not compiled with nvJPEG support" );
8
10
}
11
+ } // namespace image
12
+ } // namespace vision
9
13
#else
10
14
11
15
#include < ATen/ATen.h>
@@ -111,9 +115,10 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
111
115
CUDAJpegEncoder::CUDAJpegEncoder (const torch::Device& target_device)
112
116
: original_device{torch::kCUDA , torch::cuda::current_device ()},
113
117
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 )} {
117
122
nvjpegStatus_t status;
118
123
status = nvjpegCreateSimple (&nvjpeg_handle);
119
124
TORCH_CHECK (
You can’t perform that action at this time.
0 commit comments