Use PyTorch as logits transpose for ONNX support (openai#141)

MiscellaneousStuff · Sep 26, 2022 · 9c8183a · 9c8183a
1 parent 2037b65
commit 9c8183a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/whisper/model.py b/whisper/model.py
@@ -189,7 +189,7 @@ def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = None):
             x = block(x, xa, mask=self.mask, kv_cache=kv_cache)
 
         x = self.ln(x)
-        logits = (x @ self.token_embedding.weight.to(x.dtype).T).float()
+        logits = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float()
 
         return logits