From 2377e5f6ee2a7b7d9e22b449570d8850979b5eb0 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 13 Jul 2020 15:47:46 -0700 Subject: [PATCH] FP16 EMA bug fix --- utils/torch_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 49d08ba7b4b7..06d044779410 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -195,8 +195,8 @@ class ModelEMA: def __init__(self, model, decay=0.9999, updates=0): # Create EMA self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA - if next(model.parameters()).device.type != 'cpu': - self.ema.half() # FP16 EMA + # if next(model.parameters()).device.type != 'cpu': + # self.ema.half() # FP16 EMA self.updates = updates # number of EMA updates self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) for p in self.ema.parameters():