[FRONTEND]Darknet support batch size for yolo (apache#5688)

Fix the issue reported in https://discuss.tvm.ai/t/yolov3-tiny-batch-input-test-failed/6796
neo-ai · Jun 18, 2020 · bd2f1a6 · bd2f1a6
1 parent 50e6724
commit bd2f1a6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tvm/relay/frontend/darknet.py b/python/tvm/relay/frontend/darknet.py
@@ -637,12 +637,12 @@ def _get_darknet_attrs(self, layer, layer_num):
             attr.update({'coords' : layer.coords})
             attr.update({'background' : layer.background})
             attr.update({'softmax' : layer.softmax})
-            attr.update({'shape' : (1, layer.c, layer.h, layer.w)})
+            attr.update({'shape' : (-1, layer.c, layer.h, layer.w)})
 
         elif LAYERTYPE.YOLO == layer_type:
             attr.update({'n' : layer.n})
             attr.update({'classes' : layer.classes})
-            attr.update({'shape' : (1, layer.c, layer.h, layer.w)})
+            attr.update({'shape' : (-1, layer.c, layer.h, layer.w)})
 
         elif LAYERTYPE.UPSAMPLE == layer_type:
             attr.update({'scale' : layer.stride})