diff --git a/models.py b/models.py
index 330f4621..528ede98 100755
--- a/models.py
+++ b/models.py
@@ -167,7 +167,9 @@ class YOLOLayer(nn.Module):
                 create_grids(self, img_size, (nx, ny), p.device, p.dtype)
 
         # p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85)  # (bs, anchors, grid, grid, classes + xywh)
-        p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous()  # prediction
+        # p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous()  # prediction
+        # https://discuss.pytorch.org/t/in-pytorch-0-4-is-it-recommended-to-use-reshape-than-view-when-it-is-possible/17034
+        p = p.reshape(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2)  # prediction
 
         if self.training:
             return p