bug fix

ud.py

@@ -85,7 +85,7 @@ class UniversalDependenciesCausalPipeline(BellmanFordTokenClassificationPipeline
     m=torch.stack(m)
     k=list(range(len(d)+1))
     with torch.no_grad():
-      e=self.model(inputs_embeds=torch.stack([m[k+list(range(i,len(d)))+[-1]*i,:] for i in range(len(d))])).logits[:,-len(d):,:].numpy()
+      e=self.model(inputs_embeds=torch.stack([m[k+list(range(i,len(d)))+[-1]*i,:] for i in range(len(d))]).to(self.device)).logits[:,-len(d):,:].numpy()
     for i in range(len(d)):
       for j in range(i):
         e[-j-1,-i-1],e[-i-1,-j-1]=e[-i-1,i-j]+self.left_arc,e[-i-1,i-j]+self.right_arc
@@ -118,4 +118,7 @@ class UniversalDependenciesCausalPipeline(BellmanFordTokenClassificationPipeline
     z=matrix-numpy.nanmax(matrix,axis=0)
     m=numpy.block([[z[x,:][:,x],numpy.nanmax(z[x,:][:,y],axis=1).reshape(len(x),1)],[numpy.nanmax(z[y,:][:,x],axis=0),numpy.nanmax(z[y,y])]])
     k=[j if i==len(x) else x[j] if j<len(x) else y[numpy.nanargmax(z[y,x[i]])] for i,j in enumerate(self.chu_liu_edmonds(m))]
-
+    h=[j if i in y else k[x.index(i)] for i,j in enumerate(h)]
+    i=y[numpy.nanargmax(z[x[k[-1]],y] if k[-1]<len(x) else z[y,y])]
+    h[i]=x[k[-1]] if k[-1]<len(x) else i
+    return h