(config, ckpt, gpu, eval_mode)
| 226 | |
| 227 | |
| 228 | def load_model(config, ckpt, gpu, eval_mode): |
| 229 | if ckpt: |
| 230 | print(f"Loading model from {ckpt}") |
| 231 | pl_sd = torch.load(ckpt, map_location="cpu") |
| 232 | global_step = pl_sd["global_step"] |
| 233 | else: |
| 234 | pl_sd = {"state_dict": None} |
| 235 | global_step = None |
| 236 | model = load_model_from_config(config.model, |
| 237 | pl_sd["state_dict"]) |
| 238 | |
| 239 | return model, global_step |
| 240 | |
| 241 | |
| 242 | if __name__ == "__main__": |
no test coverage detected