There was an error while loading. Please reload this page.
1 parent 1d8f928 commit 4e85793Copy full SHA for 4e85793
1 file changed
mmcv/ops/csrc/pytorch/npu/focal_loss_npu.cpp
@@ -17,7 +17,9 @@ void sigmoid_focal_loss_forward_npu(Tensor input, Tensor target, Tensor weight,
17
int64_t weight_size = weight.size(0);
18
at::Tensor weight_y = at::ones_like(input);
19
if (weight_size > 0) {
20
- weight_y = at::broadcast_to(weight, input.sizes());
+ at::Tensor weight_selected = weight.gather(0, target);
21
+ weight_selected = weight_selected.unsqueeze(1);
22
+ weight_y = weight_selected.expand_as(input);
23
}
24
OpCommand cmd;
25
string reduction = "none";
0 commit comments