Removed non-persistent buffers from state_dict [skip ci]

ankane · Aug 6, 2024 · b509c33 · b509c33
1 parent 613e6b8
commit b509c33
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lib/torch/nn/module.rb b/lib/torch/nn/module.rb
@@ -407,7 +407,10 @@ def save_to_state_dict(destination, prefix: "")
           destination[prefix + k] = v
         end
         named_buffers.each do |k, v|
-          destination[prefix + k] = v
+          # TODO exclude v.nil?
+          if !@non_persistent_buffers_set.include?(k)
+            destination[prefix + k] = v
+          end
         end
       end