Fix datatype issue

yzhangcs · yzhangcs · commit fe4ae6f5d102 · 2021-08-02T13:52:47.000+08:00
diff --git a/supar/modules/transformer.py b/supar/modules/transformer.py
@@ -59,7 +59,7 @@ def reset_parameters(self):
         self.embed.weight.copy_(w)
 
     def forward(self, x):
-        pos = x.new_tensor(range(x.shape[1]))
+        pos = x.new_tensor(range(x.shape[1])).long()
         offset = sum(divmod(self.embed.weight.shape[0], 2))
         return self.embed(pos - pos.unsqueeze(-1) + offset)
 
@@ -179,7 +179,7 @@ def __init__(self, n_heads, n_model, n_embed, dropout=0.1):
         self.n_embed = n_embed
         self.scale = n_embed**0.5
 
-        self.pos_embed = SinusoidRelativePositionalEmbedding()
+        self.pos_embed = RelativePositionalEmbedding(n_model=n_embed)
         self.wq = nn.Parameter(torch.zeros(n_model, n_embed, n_heads))
         self.wk = nn.Parameter(torch.zeros(n_model, n_embed, n_heads))
         self.wv = nn.Parameter(torch.zeros(n_model, n_embed, n_heads))
