Mathis Petrovich commited on
Commit
007d500
1 Parent(s): 7204931

remove tuple in type annotation

Browse files
Files changed (1) hide show
  1. model.py +2 -3
model.py CHANGED
@@ -1,4 +1,4 @@
1
- from typing import List, Union
2
  import torch.nn as nn
3
  import os
4
 
@@ -64,8 +64,7 @@ class TMR_textencoder(nn.Module):
64
  )
65
 
66
  def get_last_hidden_state(self, texts: List[str],
67
- return_mask: bool = False
68
- ) -> Union[Tensor, tuple[Tensor, Tensor]]:
69
  encoded_inputs = self.tokenizer(texts, return_tensors="pt", padding=True)
70
  output = self.text_model(**encoded_inputs.to(self.text_model.device))
71
  if not return_mask:
 
1
+ from typing import List
2
  import torch.nn as nn
3
  import os
4
 
 
64
  )
65
 
66
  def get_last_hidden_state(self, texts: List[str],
67
+ return_mask: bool = False):
 
68
  encoded_inputs = self.tokenizer(texts, return_tensors="pt", padding=True)
69
  output = self.text_model(**encoded_inputs.to(self.text_model.device))
70
  if not return_mask: