InstColorization icon indicating copy to clipboard operation
InstColorization copied to clipboard

Support CPU Test.

Open delldu opened this issue 4 years ago • 1 comments

Fix Bug: model load wrong parameters and forward with wrong data type when CPU test.

delldu avatar Nov 26 '20 03:11 delldu

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)

sysuzyq avatar Apr 01 '21 09:04 sysuzyq