1 #pragma once 2 3 #ifdef USE_C10D_NCCL 4 5 #include <ATen/ATen.h> 6 #include <c10/cuda/CUDAStream.h> 7 8 namespace c10d { 9 10 // Check for NaNs in a tensor on a given stream. If any are found, throw a 11 // device-side error. 12 void checkForNan(const at::Tensor& tensor, at::cuda::CUDAStream& stream); 13 14 } // namespace c10d 15 16 #endif // USE_C10D_NCCL 17