Fix a mistake in the previous commit

This commit is contained in:
Gnome Ann 2022-02-02 23:04:59 -05:00
parent 78f52063c7
commit 4904af6adc
1 changed files with 1 additions and 1 deletions

View File

@ -755,7 +755,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly", "TPUMeshTransforme
if(hasattr(self, "transformer")): if(hasattr(self, "transformer")):
inputs_embeds = self.transformer.wte(input_ids) inputs_embeds = self.transformer.wte(input_ids)
else: else:
inputs_embeds = self.model.embed_tokens(input_ids) * self.embed_scale inputs_embeds = self.model.embed_tokens(input_ids) * self.model.embed_scale
if(vars.sp is not None): if(vars.sp is not None):
vars.sp = vars.sp.to(inputs_embeds.dtype).to(inputs_embeds.device) vars.sp = vars.sp.to(inputs_embeds.dtype).to(inputs_embeds.device)
inputs_embeds = torch.where( inputs_embeds = torch.where(