Fix handling of nonpersistent buffers (#2604)

Summary:
Pull Request resolved: #2604
bypass-github-pytorch-ci-checks

Reviewed By: cccclai

Differential Revision: D55268323

fbshipit-source-id: 40b7928e204103d8d67482a7f3332b5910b13de8

Author: angelayi

exir/lowered_backend_module.py

@@ -483,10 +483,22 @@ def _get_new_signature(  # noqa: C901
             elif is_tagged:
                 input_specs.append(orig_input_spec)
 
-                if orig_input_spec.kind in (InputKind.PARAMETER, InputKind.BUFFER):
+                if orig_input_spec.kind == InputKind.PARAMETER:
                     new_state_dict[orig_input_spec.target] = (
                         original_program.state_dict[orig_input_spec.target]
                     )
+                elif (
+                    orig_input_spec.kind == InputKind.BUFFER
+                    and orig_input_spec.persistent
+                ):
+                    new_state_dict[orig_input_spec.target] = (
+                        original_program.state_dict[orig_input_spec.target]
+                    )
+                elif orig_input_spec.kind == InputKind.BUFFER:
+                    assert not orig_input_spec.persistent
+                    new_constants[orig_input_spec.target] = original_program.constants[
+                        orig_input_spec.target
+                    ]
                 elif orig_input_spec.kind in (
                     InputKind.CONSTANT_TENSOR,
                     InputKind.CUSTOM_OBJ,