From 9086caf0bb8b31ec59d64c9c83669d26b20a9b38 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 16 Feb 2019 14:47:16 +0100 Subject: [PATCH] updates --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models.py b/models.py index d6a50eeb..181950f9 100755 --- a/models.py +++ b/models.py @@ -146,7 +146,7 @@ class YOLOLayer(nn.Module): def forward(self, p, targets=None, var=None): FT = torch.cuda.FloatTensor if p.is_cuda else torch.FloatTensor - bs = 1 if ONNX_EXPORT else p.shape[0] # batch size + bs = p.shape[0] # batch size nG = self.nG # number of grid points if p.is_cuda and not self.weights.is_cuda: