allow cudnn for fp16 batch norm (#4021)

ngimel · soumith · commit d7bd3b9acf1c · 2018-02-06T20:53:04.000-08:00
diff --git a/torch/csrc/autograd/functions/batch_normalization.cpp b/torch/csrc/autograd/functions/batch_normalization.cpp
@@ -55,7 +55,8 @@ auto BatchNormForward::apply(const variable_list& inputs) -> variable_list {
   bool use_cudnn = false;
 #ifdef WITH_CUDNN
   use_cudnn = (input.type().isCuda()
-               && input.type().scalarType() != at::kHalf
+               && (input.type().scalarType() != at::kHalf
+               || weight.type().scalarType() == at::kFloat)
                && weight.defined() && bias.defined()
                && input.size(0) <= 131070
                && cudnn_enabled && CUDNN_VERSION >= 5110L);
@@ -115,7 +116,8 @@ auto BatchNormBackward::apply(const variable_list& grad_outputs) -> variable_lis
   bool use_cudnn = false;
 #ifdef WITH_CUDNN
   use_cudnn = (input.type().backend() == at::kCUDA
-               && input.type().scalarType() != at::kHalf
+               && (input.type().scalarType() != at::kHalf
+               || weight.type().scalarType() == at::kFloat)
                && weight.defined() && bias.defined() && training
                && input.size(0) <= 131070
                && cudnn_enabled && CUDNN_VERSION >= 5110L);
