Home
last modified time | relevance | path

Searched defs:num_of_segments (Results 1 – 1 of 1) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DEmbeddingBackwardKernel.cu49 int64_t num_of_segments = *num_of_segments_ptr; in krn_partials_per_segment() local
67 int64_t num_of_segments = *num_of_segments_ptr; in krn_partial_segment_offset() local
90 int64_t num_of_segments = *num_of_segments_ptr; in compute_grad_weight_bags() local
135 int64_t num_of_segments = *num_of_segments_ptr; in compute_grad_weight() local
168 int64_t num_of_segments = *num_of_segments_ptr; in sum_and_scatter() local
194 int64_t num_of_segments = *num_of_segments_ptr; in compute_num_of_partial_segments() local
200 …ite_num_of_segments_for_legacy_thrust_path(int64_t *num_of_segments_ptr, int64_t num_of_segments) { in write_num_of_segments_for_legacy_thrust_path()
239 …int64_t num_of_segments = embedding_backward_cuda_kernel_unique_by_key<index_t>(sorted_indices, se… in embedding_backward_cuda_kernel() local