From 7c47d3e663375743e28bfd8863af1091f3b7c8b1 Mon Sep 17 00:00:00 2001
From: Benoit Steiner <benoit.steiner.goog@gmail.com>
Date: Thu, 7 Apr 2016 22:50:25 -0700
Subject: Fixed the type casting benchmarks for fp16

---
 bench/tensors/tensor_benchmarks.h | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

(limited to 'bench')

diff --git a/bench/tensors/tensor_benchmarks.h b/bench/tensors/tensor_benchmarks.h
index 16b388abf..90b9bc741 100644
--- a/bench/tensors/tensor_benchmarks.h
+++ b/bench/tensors/tensor_benchmarks.h
@@ -46,8 +46,13 @@ template <typename Device, typename T> class BenchmarkSuite {
   void typeCasting(int num_iters) {
     eigen_assert(m_ == n_);
     Eigen::array<TensorIndex, 2> sizes;
-    sizes[0] = m_;
-    sizes[1] = k_;
+    if (sizeof(T) >= sizeof(int)) {
+      sizes[0] = m_;
+      sizes[1] = k_;
+    } else {
+      sizes[0] = m_ * sizeof(T) / sizeof(int);
+      sizes[1] = k_ * sizeof(T) / sizeof(int);
+    }
     const TensorMap<Tensor<int, 2, 0, TensorIndex>, Eigen::Aligned> A((int*)a_, sizes);
     TensorMap<Tensor<T, 2, 0, TensorIndex>, Eigen::Aligned> B(b_, sizes);
 
-- 
cgit v1.2.3