Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add model unittest #8

Merged
merged 9 commits into from
Dec 2, 2020
Merged
Prev Previous commit
Next Next commit
Add type annotations
zhiqwang committed Dec 2, 2020

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
commit c20ef26afa0252ff6655afadc1ee1e71e7a155eb
2 changes: 1 addition & 1 deletion models/box_head.py
Original file line number Diff line number Diff line change
@@ -16,7 +16,7 @@ def _sum(x: List[Tensor]) -> Tensor:


class YoloHead(nn.Module):
def __init__(self, in_channels, num_anchors, num_classes): # detection layer
def __init__(self, in_channels: List[int], num_anchors: int, num_classes: int):
super().__init__()
self.num_anchors = num_anchors # anchors
self.num_outputs = num_classes + 5 # number of outputs per anchor