| # Force weights_only=False everywhere in lightning cloud_io.py | |
| # PyTorch 2.6+ treats weights_only=None as True, breaking UniRig checkpoints | |
| import glob, re | |
| for pattern in [ | |
| '/root/miniconda/envs/unirig/lib/python*/site-packages/lightning/fabric/utilities/cloud_io.py', | |
| '/root/miniconda/envs/unirig/lib/python*/site-packages/lightning/fabric/plugins/io/torch_io.py', | |
| ]: | |
| for path in glob.glob(pattern): | |
| with open(path) as f: | |
| src = f.read() | |
| # Replace any torch.load call that passes weights_only as a variable | |
| # with a hardcoded False | |
| new_src = re.sub( | |
| r'weights_only\s*=\s*weights_only', | |
| 'weights_only=False', | |
| src | |
| ) | |
| # Also catch weights_only=None | |
| new_src = re.sub( | |
| r'weights_only\s*=\s*None(?!\s*\))', | |
| 'weights_only=False', | |
| new_src | |
| ) | |
| if new_src != src: | |
| with open(path, 'w') as f: | |
| f.write(new_src) | |
| print(f'Patched {path}') | |
| else: | |
| print(f'No change needed: {path}') | |