fix dot(csr.T, dns)=dns can't be called on cpu and gpu (#11087)

apache · May 31, 2018 · 8be4b8e · 8be4b8e
1 parent 9feecce
commit 8be4b8e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/operator/tensor/dot-inl.h b/src/operator/tensor/dot-inl.h
@@ -246,6 +246,10 @@ inline bool DotForwardInferStorageType(const nnvm::NodeAttrs& attrs,
     if (target_stype == kRowSparseStorage) {
       dispatched = storage_type_assign(&out_stype, kRowSparseStorage,
                                        dispatch_mode, DispatchMode::kFComputeEx);
+    // csr.T, rsp/dns -> dns
+    } else if (target_stype == kDefaultStorage) {
+      dispatched = storage_type_assign(&out_stype, kDefaultStorage, dispatch_mode,
+                                       DispatchMode::kFComputeEx);
     }
   }
   if (!dispatched && lhs_stype == kCSRStorage && rhs_rsp_or_dns &&