FP16 test loss bug fix

5 years ago · 0b514daced
parent a1748a8d6e
commit 0b514daced
1 changed files with 1 additions and 1 deletions
--- a/test.py
+++ b/test.py
@ -104,7 +104,7 @@ def test(data,

            # Compute loss
            if training:  # if model has loss hyperparameters
-                loss += compute_loss(train_out, targets, model)[1][:3]  # GIoU, obj, cls
+                loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3]  # GIoU, obj, cls

            # Run NMS
            t = torch_utils.time_synchronized()