lib/torch/utils/data/tensor_dataset.rb in torch-rb-0.1.1 vs lib/torch/utils/data/tensor_dataset.rb in torch-rb-0.1.2

@@ -5,10 +5,14 @@
         def initialize(*tensors)
           @tensors = tensors
         end
 
         def [](index)
-          tensors.map { |t| t[index] }
+          @tensors.map { |t| t[index] }
+        end
+
+        def size
+          @tensors[0].size(0)
         end
       end
     end
   end
 end