module Torch module NN class Module def initialize @training = true end def inspect str = String.new str << "#{self.class.name}(\n" modules.each do |name, mod| str << " (#{name}): #{mod.inspect}\n" end str << ")" end def train(mode = true) @training = mode modules.each do |_, mod| mod.train(mode) end end def eval train(false) end def call(*input) forward(*input) end # modifies in-place def to(device) instance_variables.each do |name| param = instance_variable_get(name) if param.is_a?(Parameter) instance_variable_set(name, Parameter.new(param.to(device))) end end modules.each do |_, mod| mod.to(device) end self end def parameters params = [] instance_variables.each do |name| param = instance_variable_get(name) params << param if param.is_a?(Parameter) end params + modules.flat_map { |_, mod| mod.parameters } end def zero_grad parameters.each do |param| if param.grad param.grad.detach! param.grad.zero! end end end def method_missing(method, *args, &block) modules[method.to_s] || super end def respond_to?(method, include_private = false) modules.key?(method.to_s) || super end private def modules modules = {} instance_variables.each do |name| mod = instance_variable_get(name) modules[name[1..-1]] = mod if mod.is_a?(Module) end modules end end end end