lib/dnn/core/model.rb in ruby-dnn-0.8.4 vs lib/dnn/core/model.rb in ruby-dnn-0.8.5

@@ -214,10 +214,16 @@
       else
         layer_class, index = args
         @layers.select { |layer| layer.is_a?(layer_class) }[index]
       end
     end
+
+    def get_all_layers
+      @layers.map { |layer|
+        layer.is_a?(Model) ? layer.get_all_layers : layer
+      }.flatten
+    end
   
     def forward(x, training)
       @training = training
       @layers.each do |layer|
         x = if layer.is_a?(Layers::Layer)
@@ -289,10 +295,10 @@
       end
     end
 
     def layers_shape_check
       @layers.each.with_index do |layer, i|
-        prev_shape = layer.prev_layer.shape
+        prev_shape = layer.is_a?(Layers::Layer) ? layer.prev_layer.shape : layer.layers[-1]
         if layer.is_a?(Layers::Dense)
           if prev_shape.length != 1
             raise DNN_ShapeError.new("layer index(#{i}) Dense:  The shape of the previous layer is #{prev_shape}. The shape of the previous layer must be 1 dimensional.")
           end
         elsif layer.is_a?(Layers::Conv2D) || layer.is_a?(Layers::MaxPool2D)