Fixed a bug

Xeadriel · Nov 2, 2022 · 48f71dc · 48f71dc
1 parent 0a516b9
commit 48f71dc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 15 deletions.
diff --git a/firelang/models/_fireword.py b/firelang/models/_fireword.py
@@ -9,7 +9,7 @@
 from firelang.measure import Measure
 from firelang.function import Functional
 from firelang.stack import StackingSlicing
-from firelang.measure import DiracMixture, metrics
+from firelang.measure import DiracMixture
 from firelang.utils.timer import Timer, elapsed
 from firelang.utils.optim import Loss
 
@@ -139,7 +139,9 @@ def field(
         return outputs
 
     def loss_skipgram(
-        self, pairs: Tensor, labels: Tensor, args: Namespace = Namespace()
+        self,
+        pairs: Tensor,
+        labels: Tensor,
     ) -> Loss:
         """Noise contrastive estimation loss for the SkipGram task.
 
@@ -162,19 +164,6 @@ def loss_skipgram(
         )
         loss.add("sim", loss_sim)
 
-        if hasattr(args, "sinkhorn_weight") and args.sinkhorn_weight > 0.0:
-            s = metrics.sinkhorn(
-                measure1,
-                measure2,
-                reg=args.sinkhorn_reg,
-                max_iter=args.sinkhorn_max_iter,
-                p=args.sinkhorn_p,
-                tau=args.sinkhorn_tau,
-                stop_threshold=args.sinkhorn_stop_threshold,
-            )  # (n,)
-            s[~labels] = -s[~labels]
-            loss.add("sinkhorn", s * args.sinkhorn_weight)
-
         return loss
 
 

diff --git a/scripts/train.py b/scripts/train.py
@@ -1,3 +1,4 @@
+from __future__ import annotations
 from typing import List
 import argparse
 import os