(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len)
| 183 | |
| 184 | class Tokens: |
| 185 | def __init__(self, input_ids, cu_input_ids, attention_mask, cu_seqlens, max_seq_len) -> None: |
| 186 | self.input_ids = input_ids |
| 187 | self.attention_mask = attention_mask |
| 188 | self.cu_input_ids = cu_input_ids |
| 189 | self.cu_seqlens = cu_seqlens |
| 190 | self.max_seq_len = max_seq_len |
| 191 | def to(self, device): |
| 192 | self.input_ids = self.input_ids.to(device) |
| 193 | self.attention_mask = self.attention_mask.to(device) |