ldy/bge_finetune
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revised saving

Browse Source
This commit is contained in:
ldy 2025-07-23 15:17:14 +08:00
parent 8a7a011cb1
commit bace84f9e9
2 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tests/test_installation.py
View File

@ -525,7 +525,12 @@ class ComprehensiveTestSuite:
# Test saving # Test saving
state_dict = { state_dict = {
'model_state_dict': torch.randn(10, 10), 'model_state_dict': {
'layer1.weight': torch.randn(10, 5),
'layer1.bias': torch.randn(10),
'layer2.weight': torch.randn(3, 10),
'layer2.bias': torch.randn(3)
},
'optimizer_state_dict': {'lr': 0.001}, 'optimizer_state_dict': {'lr': 0.001},
'step': 100 'step': 100
} }

2
utils/checkpoint.py
View File

@ -86,7 +86,7 @@ class CheckpointManager:
model_path = os.path.join(temp_dir, 'pytorch_model.bin') model_path = os.path.join(temp_dir, 'pytorch_model.bin')
# Add extra safety check for model state dict # Add extra safety check for model state dict
model_state = state_dict['model_state_dict'] model_state = state_dict['model_state_dict']
if not model_state: if model_state is None or (isinstance(model_state, dict) and len(model_state) == 0):
raise RuntimeError("Model state dict is empty") raise RuntimeError("Model state dict is empty")
self._safe_torch_save(model_state, model_path) self._safe_torch_save(model_state, model_path)