round scale

apache · Jul 31, 2018 · d1a6fdb · d1a6fdb
1 parent 38956fd
commit d1a6fdb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/cpp/operator/mkldnn.cc b/tests/cpp/operator/mkldnn.cc
@@ -629,7 +629,7 @@ std::vector<NDArrayAttrs> GetTestInputArrays(bool rand = false,
       continue;
 
     for (int dim = 0; dim < scale.size(); dim++)
-      shape[dim] = shape[dim] * scale[dim];
+      shape[dim] = static_cast<int>(round(shape[dim] * scale[dim]));
 
     // Type 1.
     NDArray arr(shape, Context());
@@ -702,7 +702,7 @@ std::vector<NDArrayAttrs> GetTestOutputArrays(
   TShape shape = shp;
 
   for (int dim = 0; dim < scale.size(); dim++)
-    shape[dim] = static_cast<int>(shape[dim] * scale[dim]);
+    shape[dim] = static_cast<int>(round(shape[dim] * scale[dim]));
 
   std::vector<NDArrayAttrs> in_arrs;
   std::string desc;