tf: layers: yolo_layer: correct docs

hhk7734 · Feb 27, 2021 · 3d31292 · 3d31292
1 parent 37fa5db
commit 3d31292
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/py_src/yolov4/tf/layers/yolo_layer.py b/py_src/yolov4/tf/layers/yolo_layer.py
@@ -38,9 +38,9 @@ def __init__(self, metalayer, metanet):
         @tf.function
         def _coords_0(x, training):
             """
-            @param `x`: Dim(height, width, height, channels)
+            @param `x`: Dim(batch, height, width, channels)
 
-            @return: Dim(height, width, height, channels)
+            @return: Dim(batch, height, width, channels)
                 xy: logistic, scale
                 wh: raw or exp(training)
                 oc: logistic
@@ -82,9 +82,9 @@ def _coords_0(x, training):
         @tf.function
         def _coords_1(x, training):
             """
-            @param `x`: Dim(height, width, height, channels)
+            @param `x`: Dim(batch, height, width, channels)
 
-            @return: Dim(height, width, height, channels)
+            @return: Dim(batch, height, width, channels)
                 xy: scale
                 wh: raw or pow(training)
                 oc: raw
@@ -128,8 +128,8 @@ def _coords_1(x, training):
 
     def call(self, x, training=False):
         """
-        @param `x`: Dim(height, width, height, channels)
+        @param `x`: Dim(batch, height, width, channels)
 
-        @return: Dim(height, width, height, channels)
+        @return: Dim(batch, height, width, channels)
         """
         return self._yolo_function(x, training)