Make sure to cast vars.sp to the correct dtype

This commit is contained in:
Gnome Ann 2021-10-28 13:22:07 -04:00
parent 1556bd32a5
commit a1ae11630a
1 changed files with 1 additions and 1 deletions

View File

@ -526,7 +526,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
inputs_embeds = self.transformer.wte(input_ids) inputs_embeds = self.transformer.wte(input_ids)
input_ids -= self.config.vocab_size input_ids -= self.config.vocab_size
if(vars.sp is not None): if(vars.sp is not None):
vars.sp = vars.sp.to(inputs_embeds.device) vars.sp = vars.sp.to(inputs_embeds.dtype).to(inputs_embeds.device)
inputs_embeds = torch.where( inputs_embeds = torch.where(
(input_ids >= 0)[:, :, None], (input_ids >= 0)[:, :, None],
vars.sp[input_ids.clamp(min=0)], vars.sp[input_ids.clamp(min=0)],