aboutsummaryrefslogtreecommitdiffhomepage
path: root/bench/tensors/tensor_benchmarks.h
diff options
context:
space:
mode:
Diffstat (limited to 'bench/tensors/tensor_benchmarks.h')
-rw-r--r--bench/tensors/tensor_benchmarks.h9
1 files changed, 7 insertions, 2 deletions
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);