glenn-jocher commited on
Commit
0b514da
1 Parent(s): a1748a8

FP16 test loss bug fix

Browse files
Files changed (1) hide show
  1. test.py +1 -1
test.py CHANGED
@@ -104,7 +104,7 @@ def test(data,
104
 
105
  # Compute loss
106
  if training: # if model has loss hyperparameters
107
- loss += compute_loss(train_out, targets, model)[1][:3] # GIoU, obj, cls
108
 
109
  # Run NMS
110
  t = torch_utils.time_synchronized()
 
104
 
105
  # Compute loss
106
  if training: # if model has loss hyperparameters
107
+ loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls
108
 
109
  # Run NMS
110
  t = torch_utils.time_synchronized()