1 #define TORCH_ASSERT_NO_OPERATORS 2 #include <ATen/cuda/CUDAGeneratorImpl.h> 3 #include <ATen/native/UnaryOps.h> 4 #include <ATen/native/cuda/DistributionTemplates.h> 5 6 namespace at::native { 7 geometric_kernel(TensorIteratorBase & iter,double p_,std::optional<Generator> gen)8void geometric_kernel(TensorIteratorBase& iter, double p_, std::optional<Generator> gen) { 9 auto generator = get_generator_or_default<CUDAGeneratorImpl>(gen, cuda::detail::getDefaultCUDAGenerator()); 10 at::native::templates::cuda::geometric_kernel(iter, p_, generator); 11 } 12 13 REGISTER_DISPATCH(geometric_stub, &geometric_kernel); 14 15 } // namespace at::native 16