From 4c7411617cb1fdb697dc3023765369a6389b6b31 Mon Sep 17 00:00:00 2001
From: Yichao Zhou <broken.zhou@gmail.com>
Date: Wed, 24 Jul 2019 01:38:27 -0700
Subject: [PATCH] minor improvement

---
 train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train.py b/train.py
index 8f95196..44d4731 100755
--- 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, [])),