lib/torch/utils/data/tensor_dataset.rb in torch-rb-0.2.3 vs lib/torch/utils/data/tensor_dataset.rb in torch-rb-0.2.4

@@ -1,8 +1,8 @@
 module Torch
   module Utils
     module Data
-      class TensorDataset
+      class TensorDataset < Dataset
         def initialize(*tensors)
           unless tensors.all? { |t| t.size(0) == tensors[0].size(0) }
             raise Error, "Tensors must all have same dim 0 size"
           end
           @tensors = tensors