minor

trytorch/train_model_baseline.py

@@ -43,8 +43,9 @@ class LSTM(nn.Module):
 
 class SentenceBatch():
 
-        def __init__(self, sentence_ids, tok_ids, tok_types, tok_masks, labels):
+        def __init__(self, sentence_ids, tokens, tok_ids, tok_types, tok_masks, labels):
                 self.sentence_ids = sentence_ids
+                self.tokens = tokens
                 self.tok_ids = pad_sequence(tok_ids, batch_first=True)
                 self.tok_types = pad_sequence(tok_types, batch_first=True)
                 self.tok_masks = pad_sequence(tok_masks, batch_first=True)
@@ -120,7 +121,7 @@ def collate_batch(batch):
         tok_types = [make_tok_types(l) for l in lengths]
         tok_masks = [make_tok_masks(l) for l in lengths]
         
-        return SentenceBatch(sentence_ids, tok_ids, tok_types, tok_masks, labels)
+        return SentenceBatch(sentence_ids, token_batch, tok_ids, tok_types, tok_masks, labels)
         
 def train(corpus, fmt):
         print(f'starting training of {corpus} in format {fmt}...')