ALLOW COMBINATION

Parry-Parry · Nov 14, 2024 · 81d985b · 81d985b
1 parent a14c74f
commit 81d985b
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/rankers/train/loss/__init__.py b/rankers/train/loss/__init__.py
@@ -160,6 +160,14 @@ def step_d(self):
             else:
                 self.t += 1
                 self.d_weight = self.d_weight * (self.t / self.T) ** 2
+
+        def step(self):
+            if self.t >= self.T:
+                pass
+            else:
+                self.t += 1
+                self.q_weight = self.q_weight * (self.t / self.T) ** 2
+                self.d_weight = self.d_weight * (self.t / self.T) ** 2
 
         @staticmethod
         def reg(reg, weight=0):