github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish applications/Chat/coati/models/base/actor.py code style (#4248)

pull/4338/head
Junming Wu 2023-07-18 10:43:52 +08:00 committed by binmakeswell
parent 915ed8bed1
commit 77c469e1ba
1 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
applications/Chat/coati/models/base/actor.py
View File

@ -21,16 +21,13 @@ class Actor(LoRAModule):
self.model = model self.model = model
self.convert_to_lora() self.convert_to_lora()
def forward(self, def forward(
input_ids: torch.LongTensor, self,
attention_mask: Optional[torch.Tensor] = None, input_ids: torch.LongTensor,
**model_kwargs, # HACK: `generate` method may pass more kwargs attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor: **model_kwargs, # HACK: `generate` method may pass more kwargs
) -> torch.Tensor:
"""Returns model output. """Returns model output.
""" """
output = self.model( output = self.model(input_ids, attention_mask=attention_mask, **model_kwargs)
input_ids,
attention_mask=attention_mask,
**model_kwargs
)
return output return output