From 41a002e798258114b05d0472864d330c266f0260 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 3 Apr 2020 12:38:08 -0700 Subject: [PATCH] grid.float() --- models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models.py b/models.py index 59dff6d8..3ead81f9 100755 --- a/models.py +++ b/models.py @@ -136,7 +136,7 @@ class YOLOLayer(nn.Module): # build xy offsets if not self.training: yv, xv = torch.meshgrid([torch.arange(self.ny, device=device), torch.arange(self.nx, device=device)]) - self.grid = torch.stack((xv, yv), 2).view((1, 1, self.ny, self.nx, 2)) + self.grid = torch.stack((xv, yv), 2).view((1, 1, self.ny, self.nx, 2)).float() if self.anchor_vec.device != device: self.anchor_vec = self.anchor_vec.to(device)