I am getting the following error in the below code main function 'train',
    agent.critic.load_state_dict(saved_state['critic_state_dict'])
  File "C:\Users\neosr\AppData\Roaming\Python\Python310\site-packages\torch\nn\modules\module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for Critic:
	Missing key(s) in state_dict: "gnn_model.gat1.att_src", "gnn_model.gat1.att_dst", "gnn_model.gat1.bias", "gnn_model.gat1.lin.weight", "gnn_model.gat2.att_src", "gnn_model.gat2.att_dst", "gnn_model.gat2.bias", "gnn_model.gat2.lin.weight", "fc.0.weight", "fc.0.bias", "fc.2.weight", "fc.2.bias". 
	Unexpected key(s) in state_dict: "network.0.weight", "network.0.bias", "network.2.weight", "network.2.bias". 


Code:
class CustomGNN(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(CustomGNN, self).__init__()
        self.gat1 = GATConv(in_channels, 8, heads=8, dropout=0.2) # dropout=0.6
        self.gat2 = GATConv(64, out_channels, heads=1, concat=False, dropout=0.2)  # 8*8 is 64 # dropout=0.6

    def forward(self, x, edge_index):
        x = F.dropout(x, p=0.2, training=self.training) # p=0.6
        x = F.elu(self.gat1(x, edge_index))
        # After dropout and activation
        #assert torch.equal(X1[:, 0], x[:, 0]), "Order mismatch after dropout/activation."
        x = F.dropout(x, p=0.2, training=self.training) # p=0.6
        x = self.gat2(x, edge_index)
        # Selective feature synchronization based on given mappings
        # Converting given (row, column) pairs to 0-based indexing
        sync_mappings = {
            (10, 18): (7, 18),
            (10, 19): (7, 19),
            (16, 18): (15, 18),
            (16, 19): (15, 19),
            (5, 18): (0, 18),
            (5, 19): (0, 19),            
        }
        
        original_features = x.clone()
        for target, source in sync_mappings.items():
            original_features[target[0], target[1]] = original_features[source[0], source[1]]
        
        # Only retain the transformative updates for the specified 19 features
        dynamic_updates_mask = torch.zeros_like(x, dtype=bool)
        dynamic_indices = [
            (0, 18), (0, 19), (3, 18), (3, 19), (5, 18), (5, 19), (7, 18), (7, 19),
            (10, 18), (10, 19), (13, 18), (13, 19), (15, 18), (15, 19), (16, 18),
            (16, 19), (17, 20), (18, 21), (19, 22)]       

        for row, col in dynamic_indices:
            dynamic_updates_mask[row, col] = True
        
        x = original_features * (~dynamic_updates_mask) + x * dynamic_updates_mask

        return x


class Actor(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Actor, self).__init__()
        self.gnn = gnn_model
        self.stability_head = nn.Sequential(
            nn.Linear(24, 64),  # Assuming 24 features from GNN’s processed features
            nn.ReLU(),
            nn.Linear(64, 1)    # Outputs a single stability value per node
        )
        
        # Define bounds for action space scaling
        self.bounds_low = torch.tensor([0.18e-6, 0.18e-6, 0.18e-6, 0.18e-6, 0.18e-6, 
                                        0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 0.5e-6, 
                                        15e-6, 0.1e-12, 0.8], dtype=torch.float32)
                                        
        self.bounds_high = torch.tensor([0.181e-6, 0.181e-6, 0.181e-6, 0.181e-6, 0.181e-6, 
                                         50e-6, 50e-6, 50e-6, 50e-6, 50e-6, 
                                         30e-6, 10e-12, 1.4], dtype=torch.float32)
        
    def forward(self, state):
        node_features_tensor, _, edge_index = state
        
        # Processed features from the GNN
        processed_features = self.gnn(node_features_tensor, edge_index)
        # Example assertion (add after GNN processing)
        #assert torch.all(node_features_tensor[:, 0] == processed_features[:, 0]), "Node order mismatch!"
        
        # Initialize stability predictions - assume all nodes as stable
        stability_pred = torch.ones(processed_features.size(0), device=processed_features.device)
        
        # Define the indices of nodes that actually require stability prediction
        # Given [0, 3, 5, 7, 10, 13, 15, 16] represent tunable/unstable nodes
        unstable_nodes_idx = torch.tensor([0, 3, 5, 7, 10, 13, 15, 16], dtype=torch.long, device=processed_features.device)
        
        # Fetch features for nodes requiring stability computation
        features_for_stability = processed_features[unstable_nodes_idx]
        
        # Compute stability predictions for these nodes
        stability_pred_unstable = self.stability_head(features_for_stability).squeeze(-1)
        
        # Update the overall stability predictions
        stability_pred[unstable_nodes_idx] = stability_pred_unstable
        
        # For action_indices handling, adapt based on the dynamic nature of processed_features
        action_indices = [(10, 19), (16, 19), (5, 19), (13, 19), (3, 19), (10, 18),
                          (16, 18), (5, 18), (13, 18), (3, 18), (17, 20), (18, 21), (19, 22)]
                          
        # Convert action_indices to tensor format for easier handling
        action_indices_tensor = torch.tensor(action_indices, dtype=torch.long).t()
        selected_features = processed_features[action_indices_tensor[0], action_indices_tensor[1]]
        
        # Normalize and scale the GNN outputs to the action bounds for usable actions
        normalized_gnn_output = torch.tanh(selected_features)
        action_means = self.bounds_low + (self.bounds_high - self.bounds_low) * ((normalized_gnn_output + 1) / 2)
        
        # Define action variances (assuming a simple fixed form here for demonstration)
        action_log_std = torch.zeros_like(action_means)
        action_std = torch.exp(action_log_std)

        return action_means, action_std, stability_pred

class Critic(torch.nn.Module):
    def __init__(self, gnn_model):
        super(Critic, self).__init__()
        self.gnn_model = gnn_model
        self.fc = torch.nn.Sequential(
            #torch.nn.Linear(24 * 20, 128),
            torch.nn.Linear(24 * 20, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)
        )

    def forward(self, node_features_tensor, edge_index):
        x = self.gnn_model(node_features_tensor, edge_index)
        x = x.view(-1)  # Flatten
        return self.fc(x)

class PPOAgent:
    def __init__(self, gnn_model, state_dim, action_space, lr_actor, lr_critic, gamma, gae_lambda, epsilon, policy_clip, epochs, entropy_coef):
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.epsilon = epsilon
        self.policy_clip = policy_clip
        self.epochs = epochs
        self.entropy_coef = entropy_coef  # Dynamic entropy coefficient
        
        self.actor = Actor(gnn_model)
        self.critic = Critic(gnn_model)
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.action_space = action_space  # Assume continuous
        
        self.log_dir = 'OPAMP_GNNRL_log_file25_1'
        os.makedirs(self.log_dir, exist_ok=True)
        
    def save_checkpoint(self, episode):
        torch.save(self.actor.state_dict(), os.path.join(self.log_dir, f'actor_{episode}.pt'))
        torch.save(self.critic.state_dict(), os.path.join(self.log_dir, f'critic_{episode}.pt'))

    def load_checkpoint(self, episode):
        self.actor.load_state_dict(torch.load(os.path.join(self.log_dir, f'actor_{episode}.pt')))
        self.critic.load_state_dict(torch.load(os.path.join(self.log_dir, f'critic_{episode}.pt')))

# save and load process for training process
def save_training_state(agent, episode, rewards_log, losses_log, metrics_log, filepath='training_state25_1.pt'):
    state = {
        'episode': episode,
        'actor_state_dict': agent.actor.state_dict(),
        'critic_state_dict': agent.critic.state_dict(),
        'optimizer_actor_state_dict': agent.optimizer_actor.state_dict(),
        'optimizer_critic_state_dict': agent.optimizer_critic.state_dict(),
        'rewards_log': rewards_log,
        'losses_log': losses_log,
        'metrics_log': metrics_log
    }
    torch.save(state, filepath)

def load_training_state(filepath='training_state25_1.pt'):
    if os.path.isfile(filepath):
        state = torch.load(filepath)
        return state
    else:
        return None

def train(env, agent, num_episodes, max_timesteps, batch_size):
    # Ensure simulator connection
    env.connect_to_simulator()
    
    total_samples=1000
    progress_filename='progress.pkl'
    
    start_episode = 0    
    rewards_log = []
    losses_log = []
    metrics_log = []
    episode_lengths = []
    
    dataset = None
    pretrain_params = None
    
    # Check if a previous training state exists and load it
    saved_state = load_training_state()
    if saved_state:
        start_episode = saved_state['episode']
        agent.actor.load_state_dict(saved_state['actor_state_dict'])
        agent.critic.load_state_dict(saved_state['critic_state_dict'])
        agent.optimizer_actor.load_state_dict(saved_state['optimizer_actor_state_dict'])
        agent.optimizer_critic.load_state_dict(saved_state['optimizer_critic_state_dict'])
        rewards_log = saved_state['rewards_log']
        losses_log = saved_state['losses_log']
        metrics_log = saved_state['metrics_log']
        print(f"Resuming from episode {start_episode}")