[master] get the current device of inputs, resolves #68

vacancy · Mar 30, 2022 · cf10401 · cf10401
1 parent 61b1b72
commit cf10401
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pytorch/prroi_pool/src/prroi_pooling_gpu.c b/pytorch/prroi_pool/src/prroi_pooling_gpu.c
@@ -32,7 +32,7 @@ at::Tensor prroi_pooling_forward_cuda(const at::Tensor &features, const at::Tens
         return output;
     }
 
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream(features.device().index());
     PrRoIPoolingForwardGpu(
         stream, features.data<float>(), rois.data<float>(), output.data<float>(),
         nr_channels, height, width, pooled_height, pooled_width, spatial_scale,
@@ -62,7 +62,7 @@ at::Tensor prroi_pooling_backward_cuda(
         return features_diff;
     }
 
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream(features.device().index());
     PrRoIPoolingBackwardGpu(
         stream,
         features.data<float>(), rois.data<float>(), output.data<float>(), output_diff.data<float>(),
@@ -93,7 +93,7 @@ at::Tensor prroi_pooling_coor_backward_cuda(
         return coor_diff;
     }
 
-    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream(features.device().index());
     PrRoIPoolingCoorBackwardGpu(
         stream,
         features.data<float>(), rois.data<float>(), output.data<float>(), output_diff.data<float>(),