diff --git a/models.py b/models.py
index dd77c471..18dc2b50 100755
--- a/models.py
+++ b/models.py
@@ -105,6 +105,9 @@ class YOLOLayer(nn.Module):
         self.nA = len(anchors)  # number of anchors (3)
         self.nC = nC  # number of classes (80)
         self.img_size = 0
+        # self.nG, self.stride, self.grid_xy, self.anchor_vec, self.anchor_wh = \
+        #    [], [], [], [], []
+        create_grids(self, 32, 1)
 
         if ONNX_EXPORT:  # grids must be computed in __init__
             stride = [32, 16, 8][yolo_layer]  # stride of this layer
@@ -225,6 +228,7 @@ def create_grids(self, img_size, nG, device='cpu'):
     self.anchor_wh = self.anchor_vec.view(1, self.nA, 1, 1, 2).to(device)
     self.nG = torch.FloatTensor([nG]).to(device)
 
+
 def load_darknet_weights(self, weights, cutoff=-1):
     # Parses and loads the weights stored in 'weights'
     # cutoff: save layers between 0 and cutoff (if cutoff = -1 all are saved)
diff --git a/train.py b/train.py
index 10f6332f..5e9ee0e6 100644
--- a/train.py
+++ b/train.py
@@ -12,7 +12,7 @@ def train(
         data_cfg,
         img_size=416,
         resume=False,
-        epochs=100,
+        epochs=270,
         batch_size=16,
         accumulate=1,
         multi_scale=False,