lib/torch/nn/module.rb in torch-rb-0.3.5 vs lib/torch/nn/module.rb in torch-rb-0.3.6

@@ -53,10 +53,15 @@
               # TODO should_use_set_data
               instance_variable_get(key).grad = grad_applied.requires_grad!(param.grad.requires_grad)
             end
           end
         end
-        # TODO apply to more objects
+
+        @buffers.each_key do |k|
+          buf = @buffers[k]
+          @buffers[k] = fn.call(buf) unless buf.nil?
+        end
+
         self
       end
 
       def apply(fn)
         children.each do |mod|