From d3b2cbfe7c92de888871f77967a9cfb6586dc086 Mon Sep 17 00:00:00 2001 From: James Betker Date: Thu, 11 Jun 2020 08:25:57 -0600 Subject: [PATCH] Fix loading new state dicts for RRDB --- codes/models/archs/RRDBNet_arch.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/codes/models/archs/RRDBNet_arch.py b/codes/models/archs/RRDBNet_arch.py index 1eea2624..0c8852fb 100644 --- a/codes/models/archs/RRDBNet_arch.py +++ b/codes/models/archs/RRDBNet_arch.py @@ -209,7 +209,8 @@ class RRDBNet(RRDBBase): # The parameters in self.trunk used to be in this class. To support loading legacy saves, restore them. t_state = self.trunk.state_dict() for k in t_state.keys(): - state_dict["trunk.%s" % (k,)] = state_dict.pop(k) + if k in state_dict.keys(): + state_dict["trunk.%s" % (k,)] = state_dict.pop(k) super(RRDBNet, self).load_state_dict(state_dict, strict) # Variant of RRDBNet that is "assisted" by an external pretrained image classifier whose