minor improvement

zhou13 · Jul 24, 2019 · 4c74116 · 4c74116
1 parent 437ef29
commit 4c74116
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -115,7 +115,7 @@ def main():
     if M.backbone == "stacked_hourglass":
         model = lcnn.models.hg(
             depth=M.depth,
-            head=lambda c_in, c_out: MultitaskHead(c_in, c_out),
+            head=MultitaskHead,
             num_stacks=M.num_stacks,
             num_blocks=M.num_blocks,
             num_classes=sum(sum(M.head_size, [])),