fix typo

motokimura · Dec 5, 2019 · b51eb91 · b51eb91
1 parent 4f419f0
commit b51eb91
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/models/yolo_layer.py b/models/yolo_layer.py
@@ -116,11 +116,11 @@ class (float): class index.
 
         if labels is None:  # not training
             pred[..., :4] *= self.stride
-            pred = pred.view(batchsize, -1, n_ch)  # shsape: [batch, anchor x grid_x x grid_x, n_class + 5]
+            pred = pred.view(batchsize, -1, n_ch)  # shsape: [batch, anchor x grid_y x grid_x, n_class + 5]
 
             if self.gaussian:
                 # scale objectness confidence with xywh uncertainties
-                sigma_xywh = sigma_xywh.view(batchsize, -1, 4)  # shsape: [batch, anchor x grid_x x grid_x, 4]
+                sigma_xywh = sigma_xywh.view(batchsize, -1, 4)  # shsape: [batch, anchor x grid_y x grid_x, 4]
                 sigma = sigma_xywh.mean(dim=-1)
                 pred[..., 4] *= (1.0 - sigma)
 
@@ -130,7 +130,7 @@ class (float): class index.
                 sigma_xywh[..., 2:] = torch.exp(sigma_xywh[..., 2:])
 
                 # concat pred with uncertainties
-                pred = torch.cat([pred, sigma_xywh], 2)  # shsape: [batch, anchor x grid_x x grid_x, n_class + 9]
+                pred = torch.cat([pred, sigma_xywh], 2)  # shsape: [batch, anchor x grid_y x grid_x, n_class + 9]
 
             return pred.data