From 69738279338f73f6ea1b25ee0eaa48fa41629640 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 17 Oct 2025 11:08:13 +0100 Subject: [PATCH] Check return value of cudaSetDevice --- src/torchcodec/_core/CudaDeviceInterface.cpp | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/torchcodec/_core/CudaDeviceInterface.cpp b/src/torchcodec/_core/CudaDeviceInterface.cpp index aee1ecd07..2648bf48c 100644 --- a/src/torchcodec/_core/CudaDeviceInterface.cpp +++ b/src/torchcodec/_core/CudaDeviceInterface.cpp @@ -60,12 +60,10 @@ UniqueAVBufferRef getHardwareDeviceContext(const torch::Device& device) { // Create hardware device context c10::cuda::CUDAGuard deviceGuard(device); - // Valid values for the argument to cudaSetDevice are 0 to maxDevices - 1: - // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html#group__CUDART__DEVICE_1g159587909ffa0791bbe4b40187a4c6bb - // So we ensure the deviceIndex is not negative. // We set the device because we may be called from a different thread than // the one that initialized the cuda context. - cudaSetDevice(deviceIndex); + TORCH_CHECK( + cudaSetDevice(deviceIndex) == cudaSuccess, "Failed to set CUDA device"); AVBufferRef* hardwareDeviceCtxRaw = nullptr; std::string deviceOrdinal = std::to_string(deviceIndex);