| 191 | checkpoint_freq = args.checkpoint_freq or dataset.CHECKPOINT_FREQ |
| 192 | |
| 193 | def save_checkpoint(filename): |
| 194 | if args.skip_model_save: |
| 195 | return |
| 196 | save_dict = { |
| 197 | "args": vars(args), |
| 198 | "model_input_shape": dataset.input_shape, |
| 199 | "model_num_classes": dataset.num_classes, |
| 200 | "model_num_domains": len(dataset) - len(args.test_envs), |
| 201 | "model_hparams": hparams, |
| 202 | "model_dict": algorithm.cpu().state_dict() |
| 203 | } |
| 204 | torch.save(save_dict, os.path.join(args.output_dir, filename)) |
| 205 | |
| 206 | |
| 207 | last_results_keys = None |