diff --git a/models.py b/models.py
index 9d0d7de5..5e1ac1b3 100755
--- a/models.py
+++ b/models.py
@@ -145,7 +145,7 @@ class YOLOLayer(nn.Module):
 
     def create_grids(self, ng=(13, 13), device='cpu'):
         self.nx, self.ny = ng  # x and y grid size
-        self.ng = torch.tensor(ng)
+        self.ng = torch.tensor(ng, dtype=torch.float)
 
         # build xy offsets
         if not self.training:
@@ -193,9 +193,9 @@ class YOLOLayer(nn.Module):
         elif ONNX_EXPORT:
             # Avoid broadcasting for ANE operations
             m = self.na * self.nx * self.ny
-            ng = 1 / self.ng.repeat((m, 1))
-            grid = self.grid.repeat((1, self.na, 1, 1, 1)).view(m, 2)
-            anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(m, 2) * ng
+            ng = 1. / self.ng.repeat(m, 1)
+            grid = self.grid.repeat(1, self.na, 1, 1, 1).view(m, 2)
+            anchor_wh = self.anchor_wh.repeat(1, 1, self.nx, self.ny, 1).view(m, 2) * ng
 
             p = p.view(m, self.no)
             xy = torch.sigmoid(p[:, 0:2]) + grid  # x, y