Update iwan.py to fix thuml#201

oldschool24 · Nov 27, 2023 · dcb6fd9 · dcb6fd9
1 parent b308c0f
commit dcb6fd9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tllib/reweight/iwan.py b/tllib/reweight/iwan.py
@@ -48,7 +48,7 @@ def get_importance_weight(self, feature):
             instance weight in shape :math:`(N, 1)`
         """
         weight = 1. - self.discriminator(feature)
-        weight = weight / weight.mean()
+        weight = weight / (weight.mean() + 1e-5)
         weight = weight.detach()
         return weight