fix tuple_dataset

red-data-tools · May 29, 2018 · b331a37 · b331a37
1 parent c7bceed
commit b331a37
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lib/chainer/datasets/tuple_dataset.rb b/lib/chainer/datasets/tuple_dataset.rb
@@ -21,7 +21,7 @@ def [](index)
         end
         if index.kind_of?(Enumerable)
           length = batches[0].shape[0]
-          length.times.map {|i| batches.map { |m| m[i] } }
+          length.times.map {|i| batches.map { |m| m.ndim > 1 ? m[i, false] : m[i] } }
         else
           batches
         end