Add type hints in ops/border_align.py (#2026)

pull/2029/head
ShawnHu 2022-05-30 15:33:18 +08:00 committed by GitHub
parent ea173c9f07
commit e234c183cc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions

View File

@ -85,11 +85,12 @@ class BorderAlign(nn.Module):
(e.g. top, bottom, left, right).
"""
def __init__(self, pool_size):
def __init__(self, pool_size: int):
super().__init__()
self.pool_size = pool_size
def forward(self, input, boxes):
def forward(self, input: torch.Tensor,
boxes: torch.Tensor) -> torch.Tensor:
"""
Args:
input: Features with shape [N,4C,H,W]. Channels ranged in [0,C),