updates
This commit is contained in:
parent
f728bd21d2
commit
344bea20eb
|
@ -158,7 +158,7 @@ class YOLOLayer(nn.Module):
|
||||||
|
|
||||||
tcls = tcls[mask]
|
tcls = tcls[mask]
|
||||||
if xy.is_cuda:
|
if xy.is_cuda:
|
||||||
txy, tw, th, mask, tcls = txy.cuda(), twh.cuda(), mask.cuda(), tcls.cuda()
|
txy, twth, mask, tcls = txy.cuda(), twh.cuda(), mask.cuda(), tcls.cuda()
|
||||||
|
|
||||||
# Compute losses
|
# Compute losses
|
||||||
nT = sum([len(x) for x in targets]) # number of targets
|
nT = sum([len(x) for x in targets]) # number of targets
|
||||||
|
|
Loading…
Reference in New Issue