Added prefix and recurse options to named_buffers method [skip ci]

ankane · Aug 6, 2024 · 18f2122 · 18f2122
1 parent 7baafd0
commit 18f2122
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.17.1 (unreleased)
 
 - Added `persistent` option to `register_buffers`
+- Added `prefix` and `recurse` options to `named_buffers` method
 
 ## 0.17.0 (2024-07-26)
 

diff --git a/lib/torch/nn/module.rb b/lib/torch/nn/module.rb
@@ -197,8 +197,18 @@ def buffers
         named_buffers.values
       end
 
-      def named_buffers
-        @buffers || {}
+      # TODO set recurse: true in 0.18.0
+      def named_buffers(prefix: "", recurse: false)
+        buffers = {}
+        if recurse
+          named_children.each do |name, mod|
+            buffers.merge!(mod.named_buffers(prefix: "#{prefix}#{name}.", recurse: recurse))
+          end
+        end
+        (@buffers || {}).each do |k, v|
+          buffers[[prefix, k].join] = v
+        end
+        buffers
       end
 
       def children