Spaces:
Running
Running
Mathis Petrovich
commited on
Commit
•
007d500
1
Parent(s):
7204931
remove tuple in type annotation
Browse files
model.py
CHANGED
@@ -1,4 +1,4 @@
|
|
1 |
-
from typing import List
|
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:
|