Sha256: 18b8922d42279ae1a8e504f49d0cdd4128f9e6c9ddeddb2edb47bdb5a9ad0409

Contents?: true

Size: 969 Bytes

Versions: 24

Compression:

Stored size: 969 Bytes

Contents

module Torch
  module NN
    class ModuleList < Module
      include Enumerable

      def initialize(mods = nil)
        super()

        self.concat(mods) if mods
      end

      def length
        @modules.length
      end
      alias_method :count, :length
      alias_method :size, :length

      def concat(mods)
        raise ArgumentError, "Modules should respond to #each" unless mods.respond_to?(:each)

        mods.each { |m| append m }

        self
      end

      def each(&block)
        if block_given?
          @modules.values.each(&block)
        else
          to_enum(:each)
        end
      end

      def append(mod)
        raise ArgumentError, "Provided element is not a module" unless mod.is_a?(Module)
        add_module(length.to_s, mod)
        self
      end

      def [](idx)
        if idx.is_a?(Range)
          self.class.new(@modules.values[idx])
        else
          @modules[idx.to_s]
        end
      end
    end
  end
end

Version data entries

24 entries across 24 versions & 1 rubygems

Version Path
torch-rb-0.9.1 lib/torch/nn/module_list.rb
torch-rb-0.9.0 lib/torch/nn/module_list.rb
torch-rb-0.8.3 lib/torch/nn/module_list.rb
torch-rb-0.8.2 lib/torch/nn/module_list.rb