Explorar o código

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

Michael Goin %!s(int64=2) %!d(string=hai) anos
pai
achega
9c8183a179
Modificáronse 1 ficheiros con 1 adicións e 1 borrados
  1. 1 1
      whisper/model.py

+ 1 - 1
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