torchnvjpeg icon indicating copy to clipboard operation
torchnvjpeg copied to clipboard

Multi GPU BUGS

Open WemChao opened this issue 2 years ago • 0 comments

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;

WemChao avatar Apr 10 '22 08:04 WemChao