Reset GRU context

hm-ysjiang · Jun 7, 2023 · b7546ea · b7546ea
1 parent cec2a79
commit b7546ea
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/train-selfsupervised.py b/train-selfsupervised.py
@@ -116,7 +116,7 @@ def train(args):
         if args.reset_context:
             _weight = OrderedDict()
             for key, val in checkpoint.items():
-                if '.cnet.' not in key:
+                if '.cnet.' not in key and '.update_block.gru.' not in key:
                     _weight[key] = val
             weight = _weight
         model.load_state_dict(weight, strict=(not args.allow_nonstrict))

diff --git a/train-supervised.py b/train-supervised.py
@@ -110,7 +110,7 @@ def train(args):
         if args.reset_context:
             _weight = OrderedDict()
             for key, val in checkpoint.items():
-                if '.cnet.' not in key:
+                if '.cnet.' not in key and '.update_block.gru.' not in key:
                     _weight[key] = val
             weight = _weight
         model.load_state_dict(weight, strict=(not args.allow_nonstrict))