From a1748a8d6ebee8dd10d8d4d67aff5135d2bbe4ab Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 15 Jun 2020 15:20:27 -0700 Subject: [PATCH] test during training default to FP16 --- test.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/test.py b/test.py index 60faf23..a9630fa 100644 --- a/test.py +++ b/test.py @@ -23,6 +23,7 @@ def test(data, verbose=False): # Initialize/load model and set device if model is None: + training = False device = torch_utils.select_device(opt.device, batch_size=batch_size) half = device.type != 'cpu' # half precision only supported on CUDA @@ -42,11 +43,12 @@ def test(data, if device.type != 'cpu' and torch.cuda.device_count() > 1: model = nn.DataParallel(model) - training = False else: # called by train.py - device = next(model.parameters()).device # get model device - half = False training = True + device = next(model.parameters()).device # get model device + half = device.type != 'cpu' # half precision only supported on CUDA + if half: + model.half() # to FP16 # Configure model.eval()