Spedon commited on
Commit
26cc346
·
1 Parent(s): 0c01441

Optimize decode_outputs for OpenVINO (#1535)

Browse files

Avoid ScatterND ops that will cause error in openvino model optimizer.

Files changed (1) hide show
  1. yolox/models/yolo_head.py +1 -2
yolox/models/yolo_head.py CHANGED
@@ -246,8 +246,7 @@ class YOLOXHead(nn.Module):
246
  grids = torch.cat(grids, dim=1).type(dtype)
247
  strides = torch.cat(strides, dim=1).type(dtype)
248
 
249
- outputs[..., :2] = (outputs[..., :2] + grids) * strides
250
- outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
251
  return outputs
252
 
253
  def get_losses(
 
246
  grids = torch.cat(grids, dim=1).type(dtype)
247
  strides = torch.cat(strides, dim=1).type(dtype)
248
 
249
+ outputs = torch.cat([(outputs[..., 0:2] + grids) * strides, torch.exp(outputs[..., 2:4]) * strides, outputs[..., 4:]], dim=-1)
 
250
  return outputs
251
 
252
  def get_losses(