| 42 | } |
| 43 | |
| 44 | size_t get_reserve_size( |
| 45 | const cudnnHandle_t& handle, const BNTensorDescHolder& tensor_desc) { |
| 46 | #if CUDNN_VERSION >= 7410 |
| 47 | size_t reserve_size; |
| 48 | cudnn_check(cudnnGetBatchNormalizationTrainingExReserveSpaceSize( |
| 49 | handle, tensor_desc.bn_mode, CUDNN_BATCHNORM_OPS_BN, |
| 50 | nullptr, // activationDesc |
| 51 | tensor_desc.xy_desc.desc, // xDesc |
| 52 | &reserve_size)); |
| 53 | return reserve_size; |
| 54 | #else |
| 55 | return 0; |
| 56 | #endif // CUDNN_VERSION >= 7410 |
| 57 | } |
| 58 | } // namespace batch_normalization |
| 59 | |
| 60 | using batch_normalization::BNTensorDescHolder; |
no outgoing calls