InstColorization
InstColorization copied to clipboard
Support CPU Test.
Fix Bug: model load wrong parameters and forward with wrong data type when CPU test.
https://github.com/ericsujw/InstColorization/blob/master/models/fusion_model.py#L108
GF_state_dict_temp ={}
for k,v in GF_state_dict.items():
if k.startswith('module'):
GF_state_dict_temp[k[7:]] = v
else:
GF_state_dict_temp[k] = v
GF_state_dict = GF_state_dict_temp
G_state_dict_temp ={}
for k,v in G_state_dict.items():
if k.startswith('module'):
G_state_dict_temp[k[7:]] = v
else:
G_state_dict_temp[k] = v
G_state_dict = G_state_dict_temp
GComp_state_dict_temp ={}
for k,v in GComp_state_dict.items():
if k.startswith('module'):
GComp_state_dict_temp[k[7:]] = v
else:
GComp_state_dict_temp[k] = v
GComp_state_dict = GComp_state_dict_temp
self.netGF.load_state_dict(GF_state_dict, strict=True)
self.netG.load_state_dict(G_state_dict, strict=True)
self.netGComp.load_state_dict(GComp_state_dict, strict=True)