diff options
Diffstat (limited to 'Whisper/ML/MlContext.cpp')
| -rw-r--r-- | Whisper/ML/MlContext.cpp | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/Whisper/ML/MlContext.cpp b/Whisper/ML/MlContext.cpp index 6eeae09..a226999 100644 --- a/Whisper/ML/MlContext.cpp +++ b/Whisper/ML/MlContext.cpp @@ -556,7 +556,15 @@ void MlContext::softMax( Tensor& a, float inputScale ) printSizes.print( a ); #endif constexpr uint32_t FIXED_ROW_SIZE = 1500; - eComputeShader cs = ( a.ne[ 0 ] == FIXED_ROW_SIZE ) ? eComputeShader::softMaxFixed : eComputeShader::softMax; + + eComputeShader cs; + if( a.ne[ 0 ] == FIXED_ROW_SIZE ) + cs = eComputeShader::softMaxFixed; + else if( a.ne[ 0 ] >= ( 1024 * 4 ) ) + cs = eComputeShader::softMaxLong; + else + cs = eComputeShader::softMax; + bindShader( cs ); const uint32_t nr = a.countRows(); TensorShape dummyShape; |