Fix use of deprecated tensor.type() in SegmentReduce.cpp (#64151)
authorPeter Bell <peterbell10@live.co.uk>
Tue, 14 Sep 2021 06:15:10 +0000 (23:15 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Tue, 14 Sep 2021 06:16:47 +0000 (23:16 -0700)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64151

Reviewed By: mruberry

Differential Revision: D30917268

Pulled By: ngimel

fbshipit-source-id: 63427372b651ac495d48ef552eba5fbf0e4378e9

aten/src/ATen/native/SegmentReduce.cpp

index 910fe41..8e5c846 100644 (file)
@@ -107,7 +107,7 @@ Tensor _segment_reduce_cpu_kernel(
   output_shape[axis] = segment_count;
   auto output = at::empty(output_shape, data.options());
 
-  AT_DISPATCH_INDEX_TYPES(lengths.type(), "_segment_reduce_cpu_kernel1", [&]() {
+  AT_DISPATCH_INDEX_TYPES(lengths.scalar_type(), "_segment_reduce_cpu_kernel1", [&]() {
     const auto* lengths_data = lengths.data_ptr<index_t>();
     _segment_reduce_cpu_kernel1(
         reduction, data, lengths_data, axis, initial, output, segment_count);
@@ -209,7 +209,7 @@ Tensor _segment_reduce_cpu_backward_kernel(
   auto grad_input = at::zeros({data_contig.sizes()}, grad_contig.options());
 
   AT_DISPATCH_INDEX_TYPES(
-      lengths_contig.type(), "_segment_reduce_cpu_backward_kernel1", [&]() {
+      lengths_contig.scalar_type(), "_segment_reduce_cpu_backward_kernel1", [&] {
         const auto* lengths_data = lengths_contig.data_ptr<index_t>();
         _segment_reduce_cpu_backward_kernel1(
             grad_contig,