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

2022-09-26 13:54:26 -04:00
parent 2037b65f3f
commit 9c8183a179
1 changed files with 1 additions and 1 deletions
--- a/whisper/model.py
+++ b/whisper/model.py
@@ -189,7 +189,7 @@ class TextDecoder(nn.Module):
            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