From 77c8e03c6b274d8dc4a2265db3fb3ba55664ae83 Mon Sep 17 00:00:00 2001
From: Tiance Wang <wangtiance@gmail.com>
Date: Tue, 9 May 2023 11:58:05 +0800
Subject: [PATCH] avoid negative values of running_var when using average model

---
 icefall/checkpoint.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/icefall/checkpoint.py b/icefall/checkpoint.py
index c83c56a53..d09c72328 100644
--- a/icefall/checkpoint.py
+++ b/icefall/checkpoint.py
@@ -444,6 +444,11 @@ def average_checkpoints_with_averaged_model(
         scaling_factor=weight_end,
     )
 
+    # avoid negative running variance of batchnorm layers
+    for k in avg.keys():
+        if k.split('.')[-1] == 'running_var':
+            avg[k] = torch.clip(avg[k], min=1e-2)    
+    
     return avg