fix the bug mentioned in Issue mli0603#8

DeH40 · Jan 14, 2022 · 4761da9 · 4761da9
1 parent f07b1ee
commit 4761da9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/module/sttr.py b/module/sttr.py
@@ -56,6 +56,9 @@ def _disable_batchnorm_tracking(self):
         for m in self.modules():
             if isinstance(m, nn.BatchNorm2d):
                 m.track_running_stats = False
+                # fix the bug mentioned in https://github.com/mli0603/stereo-transformer/issues/8
+                m.running_mean = None
+                m.running_var = None
 
     def _relu_inplace(self):
         """