updates
This commit is contained in:
parent
8069fc452a
commit
5e0cce771e
|
@ -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:
|
||||||
|
|
Loading…
Reference in New Issue