This commit is contained in:
Glenn Jocher 2019-03-04 15:52:51 +01:00
parent 8069fc452a
commit 5e0cce771e
1 changed files with 1 additions and 1 deletions

View File

@ -150,7 +150,7 @@ class YOLOLayer(nn.Module):
p_conf = p[..., 4] # Conf p_conf = p[..., 4] # Conf
p_cls = p[..., 5:] # Class p_cls = p[..., 5:] # Class
txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG.cpu()) txy, twh, mask, tcls = build_targets(targets, self.anchor_vec, self.nA, self.nC, nG)
tcls = tcls[mask] tcls = tcls[mask]
if xy.is_cuda: if xy.is_cuda: