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

@@ -56,10 +56,13 @@
           end
         end
 
         @buffers.each_key do |k|
           buf = @buffers[k]
-          @buffers[k] = fn.call(buf) unless buf.nil?
+          unless buf.nil?
+            @buffers[k] = fn.call(buf)
+            instance_variable_set("@#{k}", @buffers[k])
+          end
         end
 
         self
       end