multi_gpu multi_scale

This commit is contained in:
Glenn Jocher 2019-03-19 15:48:52 +02:00
parent bc989a0147
commit 973715060d
1 changed files with 2 additions and 2 deletions

View File

@ -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