"...utils/git@developer.sourcefind.cn:OpenDAS/torchaudio.git" did not exist on "0e5913d523e10e231258e1a4670431e4dab589f0"
Commit 134b5c75 authored by rusty1s's avatar rusty1s
Browse files

fix division for segment_coo_mean when input is of type long

parent 8ec6d0c6
...@@ -273,7 +273,10 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index, ...@@ -273,7 +273,10 @@ segment_coo_cuda(torch::Tensor src, torch::Tensor index,
auto count = arg_out.value(); auto count = arg_out.value();
for (int i = dim + 1; i < out.dim(); i++) for (int i = dim + 1; i < out.dim(); i++)
count = count.unsqueeze(-1); count = count.unsqueeze(-1);
out.div_(count); if (out.is_floating_point())
out.true_divide_(count);
else
out.floor_divide_(count);
} }
}); });
}); });
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment