Update decode_utils.py

decode_utils.py

@@ -1,4 +1,3 @@
-import logging
 import math
 import re
 from typing import (
@@ -14,8 +13,6 @@ import torch.nn as nn
 from tqdm import tqdm
 from transformers import PreTrainedTokenizer
 
-logger = logging.getLogger(__name__)
-
 
 def get_id_and_prob(spans, offset_map):
     prompt_length = 0
@@ -445,7 +442,7 @@ class UIEDecoder(nn.Module):
             }
 
             for k, v in batch.items():
-                batch[k] = torch.LongTensor(v, device=self.device)
+                batch[k] = torch.tensor(v, device=self.device)
 
             outputs = self(**batch)
             start_prob, end_prob = outputs[0], outputs[1]