From b0805fe32fce768f62cb62a0652c642ba670c983 Mon Sep 17 00:00:00 2001
From: Benjamin <ben@neuralmagic.com>
Date: Wed, 17 Mar 2021 17:24:58 -0400
Subject: [PATCH] add option to disable half precision in test loop

---
 test.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/test.py b/test.py
index 39e0992264ec..0664a3ddb15c 100644
--- a/test.py
+++ b/test.py
@@ -36,7 +36,8 @@ def test(data,
          save_conf=False,  # save auto-label confidences
          plots=True,
          log_imgs=0,  # number of logged images
-         compute_loss=None):
+         compute_loss=None,
+         half_precision=True):
     # Initialize/load model and set device
     training = model is not None
     if training:  # called by train.py
@@ -60,7 +61,7 @@ def test(data,
         #     model = nn.DataParallel(model)
 
     # Half
-    half = device.type != 'cpu'  # half precision only supported on CUDA
+    half = device.type != 'cpu' and half_precision  # half precision only supported on CUDA
     if half:
         model.half()