From 973715060d6f2cc23ef5b5c136e049d9374c5c01 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Tue, 19 Mar 2019 15:48:52 +0200
Subject: [PATCH] multi_gpu multi_scale

---
 models.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models.py b/models.py
index 21e871e1..ef417603 100755
--- a/models.py
+++ b/models.py
@@ -105,8 +105,8 @@ 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 = \
-            [], [], [], [], []
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        create_grids(self, 32, 1, device=device)
 
         if ONNX_EXPORT:  # grids must be computed in __init__
             stride = [32, 16, 8][yolo_layer]  # stride of this layer