Fixkeybuffer (#2512)

* fix buffers loading (awq)
OpenNMT · Nov 11, 2023 · c5c84af · c5c84af
1 parent da2fe8f
commit c5c84af
Showing 1 changed file with 6 additions and 7 deletions.
diff --git a/onmt/models/model.py b/onmt/models/model.py
@@ -46,7 +46,6 @@ def count_parameters(self, log=print):
         raise NotImplementedError
 
     def _load_param(self, name, module, param_name, param, buf_list, ckpt_t, offset):
-
         if name.split(".")[-1] in [
             "linear_keys",
             "linear_values",
@@ -73,7 +72,7 @@ def _load_param(self, name, module, param_name, param, buf_list, ckpt_t, offset)
                     row_slice_start:row_slice_end,
                 ].size()
             ), "An error in model's partition and checkpoint's slice was detected"
-            if param_name in buf_list:
+            if name + "." + param_name in buf_list:
                 module.register_buffer(
                     param_name,
                     ckpt_t[
@@ -90,7 +89,7 @@ def _load_param(self, name, module, param_name, param, buf_list, ckpt_t, offset)
             assert (
                 param.data.size() == ckpt_t[col_slice_start:col_slice_end].size()
             ), "An error in model's partition and checkpoint's slice was detected"
-            if param_name in buf_list:
+            if name + "." + param_name in buf_list:
                 module.register_buffer(
                     param_name, ckpt_t[col_slice_start:col_slice_end]
                 )
@@ -120,9 +119,9 @@ def load_state_dict(
         if device == torch.device("cpu"):
             offset = 0
         buf_list = []
+        for buf_name, buf in self.named_buffers():
+            buf_list.append(buf_name)
         for name, module in self.named_modules():
-            for buf_name, buf in module.named_buffers():
-                buf_list.append(buf_name)
             named_buf_and_param = list(module.named_buffers()) + list(
                 module.named_parameters()
             )
@@ -205,9 +204,9 @@ def load_safe_state_dict(
         if device == torch.device("cpu"):
             offset = 0
         buf_list = []
+        for buf_name, buf in self.named_buffers():
+            buf_list.append(buf_name)
         for name, module in self.named_modules():
-            for buf_name, buf in module.named_buffers():
-                buf_list.append(buf_name)
             named_buf_and_param = list(module.named_buffers()) + list(
                 module.named_parameters()
             )