| 37 | torch.cuda.empty_cache() |
| 38 | |
| 39 | def parse_cfgs(): |
| 40 | parser = argparse.ArgumentParser(description="Simple example of a training script.") |
| 41 | parser.add_argument('--config', type=str) |
| 42 | args = parser.parse_args() |
| 43 | cfgs = Config.fromfile(args.config) |
| 44 | |
| 45 | env_local_rank = int(os.environ.get("LOCAL_RANK", -1)) |
| 46 | if env_local_rank != -1 and env_local_rank != cfgs.local_rank: |
| 47 | cfgs.local_rank = env_local_rank |
| 48 | |
| 49 | # default to using the same revision for the non-ema model if not specified |
| 50 | if cfgs.non_ema_revision is None: |
| 51 | cfgs.non_ema_revision = cfgs.revision |
| 52 | |
| 53 | return cfgs |
| 54 | |
| 55 | |
| 56 | def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None): |