torchnvjpeg
torchnvjpeg copied to clipboard
Multi GPU BUGS
When this code is using a multi-card, it directly uses no CUDA Context environment switch, and pytorch uses the last environment by default, so the decoder created in front of it reports an error. Therefore I made the following changes in the C++ source code when calling the function.
torch::Tensor Decoder::decode(const std::string& data, bool stream_sync = true) {
auto orig_device = c10::cuda::current_device();
c10::cuda::set_device(device_id);
....
c10::cuda::set_device(orig_device);
return image_tensor;
std::vector<torch::Tensor> Decoder::batch_decode(const std::vector<std::string>& data_list, bool stream_sync = true) {
auto orig_device = c10::cuda::current_device();
c10::cuda::set_device(device_id);
....
c10::cuda::set_device(orig_device);
return tensor_list;