Update model loading to support checkpoint structure in Study_Models

Author: DefinetlyNotAI

tools/Study_Models.py

@@ -275,7 +275,8 @@ def save_model_summary(model_, filename="Model_Summary.txt"):
 # Load model
 model = SimpleNN(input_dim).to(DEVICE)
 if os.path.exists(MODEL_PATH):
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
+    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
+    model.load_state_dict(checkpoint["model_state_dict"])
 model.eval()
 
 # Sample input for activation visualization