mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix loading new state dicts for RRDB

Browse Source
This commit is contained in:
James Betker 2020-06-11 08:25:57 -06:00
parent 5ca53e7786
commit d3b2cbfe7c
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
codes/models/archs/RRDBNet_arch.py
View File

@ -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