Method
__init__
(self, decoder, lm_head, config, is_trt = False)
Source from the content-addressed store, hash-verified
| 119 | """ |
| 120 | |
| 121 | def __init__(self, decoder, lm_head, config, is_trt = False): |
| 122 | super().__init__() |
| 123 | self.decoder = decoder |
| 124 | self.lm_head = lm_head |
| 125 | self.config = config |
| 126 | # HuggingFace's beam search requires to set self.device. Set it to avoid application crash |
| 127 | self.device = torch.device('cuda') |
| 128 | # Use hardcoded value to extend compatibility with older HF versions. |
| 129 | self.main_input_name = "input_ids" |
| 130 | # trt uses cached and precomputed cross attention vs. framework uses the entire kv cache as output. Need to treat them differently. |
| 131 | self.is_trt = is_trt |
| 132 | |
| 133 | def prepare_inputs_for_generation( |
| 134 | self, |
Callers
nothing calls this directly
Tested by
no test coverage detected