Fixed the missing type hint.

open-mmlab · May 30, 2022 · 1cf5dde · 1cf5dde
1 parent 245a3de
commit 1cf5dde
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/mmcv/ops/focal_loss.py b/mmcv/ops/focal_loss.py
@@ -103,7 +103,8 @@ def __init__(self,
         self.register_buffer('weight', weight)
         self.reduction = reduction
 
-    def forward(self, input: torch.Tensor, target: torch.LongTensor):
+    def forward(self, input: torch.Tensor,
+                target: torch.LongTensor) -> torch.Tensor:
         return sigmoid_focal_loss(input, target, self.gamma, self.alpha,
                                   self.weight, self.reduction)
 
@@ -215,7 +216,8 @@ def __init__(self,
         self.register_buffer('weight', weight)
         self.reduction = reduction
 
-    def forward(self, input: torch.Tensor, target: torch.LongTensor):
+    def forward(self, input: torch.Tensor,
+                target: torch.LongTensor) -> torch.Tensor:
         return softmax_focal_loss(input, target, self.gamma, self.alpha,
                                   self.weight, self.reduction)