Modification to handle the case when there is no output (#2593)

Signed-off-by: Alankar Mahajan <quic_alanmaha@quicinc.com>
quic · Dec 4, 2023 · adfa587 · adfa587
1 parent 8f9e27a
commit adfa587
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py b/TrainingExtensions/torch/src/python/aimet_torch/elementwise_ops.py
@@ -370,7 +370,7 @@ def forward(self, *args) -> torch.Tensor:
 
         res = torch.Tensor(res).type(torch.int64)
         out = torch.zeros(batch_scores.shape[0] * batch_scores.shape[1] * self.max_output_boxes_per_class, 3, dtype=torch.int64)
-        indices = torch.arange(0, len(res) * len(res[0]), dtype=torch.int64)
+        indices = torch.arange(0, len(res) * 3, dtype=torch.int64)
         out.put_(indices, res)
         return out