Revert to torch.einsum

yzhangcs · yzhangcs · commit 2c373304cae7 · 2021-07-17T17:44:34.000+08:00
diff --git a/supar/modules/affine.py b/supar/modules/affine.py
@@ -2,7 +2,6 @@
 
 import torch
 import torch.nn as nn
-from opt_einsum import contract
 
 
 class Biaffine(nn.Module):
@@ -72,7 +71,7 @@ def forward(self, x, y):
         if self.bias_y:
             y = torch.cat((y, torch.ones_like(y[..., :1])), -1)
         # [batch_size, n_out, seq_len, seq_len]
-        s = contract('bxi,oij,byj->boxy', x, self.weight, y) / self.n_in ** self.scale
+        s = torch.einsum('bxi,oij,byj->boxy', x, self.weight, y) / self.n_in ** self.scale
         # remove dim 1 if n_out == 1
         s = s.squeeze(1)
 
@@ -146,9 +145,9 @@ def forward(self, x, y, z):
             x = torch.cat((x, torch.ones_like(x[..., :1])), -1)
         if self.bias_y:
             y = torch.cat((y, torch.ones_like(y[..., :1])), -1)
-        w = contract('bzk,oikj->bozij', z, self.weight)
+        w = torch.einsum('bzk,oikj->bozij', z, self.weight)
         # [batch_size, n_out, seq_len, seq_len, seq_len]
-        s = contract('bxi,bozij,byj->bozxy', x, w, y) / self.n_in ** self.scale
+        s = torch.einsum('bxi,bozij,byj->bozxy', x, w, y) / self.n_in ** self.scale
         # remove dim 1 if n_out == 1
         s = s.squeeze(1)
 
