[BUGFIX] Fix SupportDNNL for multiple inputs (#21102)

* Fix SupportDNNL for multiple inputs * Fix SupportDNNL condition for multiple inputs
apache · Jul 22, 2022 · 5e5e0e3 · 5e5e0e3
1 parent ecb5026
commit 5e5e0e3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/operator/nn/dnnl/dnnl_base-inl.h b/src/operator/nn/dnnl/dnnl_base-inl.h
@@ -184,16 +184,16 @@ enum DNNLTensorsDtypes { AllSame = 0, Mixed = 1 };
 
 template <int MinNdim, int MaxNdim, DNNLTypeMode TypeMode, DNNLTensorsDtypes MixedTensors>
 static inline bool SupportDNNL(const std::vector<NDArray>& inputs) {
-  int dtype = MixedTensors ? -1 : inputs[0].dtype();
-  if (!SupportDNNLType<TypeMode>(dtype)) {
+  if (!SupportDNNLType<TypeMode>(inputs[0].dtype())) {
     return false;
   }
+  int dtype = MixedTensors ? -1 : inputs[0].dtype();
   for (NDArray input : inputs) {
     if (dtype == -1) {
       if (!SupportDNNL<MinNdim, MaxNdim, TypeMode>(input))
         return false;
     } else {
-      if (input.dtype() != dtype && !SupportDNNLShape<MinNdim, MaxNdim>(input.shape()))
+      if (input.dtype() != dtype || !SupportDNNLShape<MinNdim, MaxNdim>(input.shape()))
         return false;
     }
   }