Revised saving

2025-07-23 15:17:14 +08:00
parent 8a7a011cb1
commit bace84f9e9
2 changed files with 7 additions and 2 deletions
--- a/tests/test_installation.py
+++ b/tests/test_installation.py
@@ -525,7 +525,12 @@ class ComprehensiveTestSuite:

                # Test saving
                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},
                    'step': 100
                }