diff --git a/summary.py b/summary.py
index 7d3441f..204c849 100644
--- a/summary.py
+++ b/summary.py
@@ -1,13 +1,13 @@
-#--------------------------------------------#
-#   该部分代码用于看网络结构
-#--------------------------------------------#
 import torch
 from torchsummary import summary
 
-from nets.yolo import YoloBody
+from nets.yolo4_tiny import YoloBody
+
+import numpy as np
+
 
 if __name__ == "__main__":
     # 需要使用device来指定网络在GPU还是CPU运行
     device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    m       = YoloBody([[3, 4, 5], [1, 2, 3]], 80).to(device)
-    summary(m, input_size=(3, 416, 416))
+    m       = YoloBody(2).to(device)
+    summary(m, input_data=(3, 416, 416))