From 41ff83a98d3030d920b184c3b87cc24270bbd611 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 12 Jul 2020 12:18:00 -0700
Subject: [PATCH] update torch_utils.py to FP16 EMA

---
 utils/torch_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/utils/torch_utils.py b/utils/torch_utils.py
index 429351a..49d08ba 100644
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -195,6 +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
         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():