(self, ckpt_path)
| 196 | return text_emb |
| 197 | |
| 198 | def load_pretrained(self, ckpt_path): |
| 199 | checkpoint = torch.load(ckpt_path, map_location="cpu") |
| 200 | state_dict = checkpoint["state_dict"] |
| 201 | for k in list(state_dict.keys()): |
| 202 | if "model" in k: |
| 203 | state_dict[k.replace("model.", "")] = state_dict.pop(k) |
| 204 | self.load_state_dict(state_dict, strict=True) |