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

Above code is my GNN network, where it takes input of the 'node_features_tensor' and 'edge_index' from state input. I need to update the below given critic network to utilize the same gnn_model which is used in the actor network, we use the same gnn_model for computing value function in the critic network, where it need the 'node_features_tensor' and 'edge_index' as inputs for computation.
But with in the function 'update_policy' present in the class 'PPOAgent' i am passing the flattened state of node_features_tensor alone in to the critic network which is not properly synchronized with the gnn_model required inputs states,
# Flattening the node_features_tensor for the critic, as your critic expects a flattened input.
flattened_state = node_features_tensor.view(-1)               
# Getting the state value prediction from the critic using the flattened state
state_value = self.critic(flattened_state).unsqueeze(0)

and also with in the main training loop 'train', i am computing the value of "values.append(agent.critic(state[0]).item())" and "next_value = agent.critic(state[0]).item() if not done else 0" by simply passing the state[0] (node_features_tensor) alone to the critic network which also does not meet the requirement of inputs in the gnn_model. so i need you to modify the below code and provide me the updated code with properly synchronizing the gnn_model requirements and critic network requirements appropriately by modifying the critic network call function properly passing the required inputs to the critic network which is utilizing the implemented gnn_model appropriately, in the updation of below code, keep remember the gradients are properly computed and synchronized with the implementation without any break.

# Flattening the node_features_tensor for the critic, as your critic expects a flattened input.
flattened_state = node_features_tensor.view(-1)               
# Getting the state value prediction from the critic using the flattened state
state_value = self.critic(flattened_state).unsqueeze(0)

"values.append(agent.critic(state[0]).item())" and "next_value = agent.critic(state[0]).item() if not done else 0"

class Critic(torch.nn.Module):
    def __init__(self):
        super(Critic, self).__init__()
        flattened_size = 20 * 24  # Number of nodes times features per node
        self.network = torch.nn.Sequential(
            torch.nn.Linear(flattened_size, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 1)
        )

    def forward(self, state):
        # Directly use state as it is already what we referred to as state[0] in evaluate
        flattened = state.view(-1)  # Flatten to [480]
        return self.network(flattened)

# Initialize Hyperparameters
in_channels = 24  # Number of input features
out_channels = 24 # Number of output actions or features
action_space = 13  # Environment’s action space

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()
        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

    def update_policy(self, prev_states, prev_actions, prev_log_probs, returns, advantages):
        device = next(self.actor.parameters()).device  # Get the device from the actor model
        advantages = torch.tensor(advantages, dtype=torch.float32, device=device)
        returns = torch.tensor(returns, dtype=torch.float32, device=device)
        prev_log_probs = torch.tensor(prev_log_probs, dtype=torch.float32, device=device)        
        
        # Initialize lists to accumulate losses
        actor_losses = []
        critic_losses = []
        
        for _ in range(self.epochs):
            
            log_probs_batch, state_values_batch, entropies_batch, stability_losses_batch = [], [], [], []
    
            for state, action in zip(prev_states, prev_actions):
                node_features_tensor, _, edge_index = state
                node_features_tensor = node_features_tensor.to(device)
                edge_index = edge_index.to(device) if edge_index is not None else None
                
                # Flattening the node_features_tensor for the critic, as your critic expects a flattened input.
                flattened_state = node_features_tensor.view(-1)  # Adjusted to flatten the state appropriately
            
                # Extracting the 24th feature for stability; assuming indexing starts at 0, so it’s 23rd in 0-indexing.
                stability_target = node_features_tensor[:, 23].unsqueeze(1)  # For the stability target
    
                action_means, action_stds, stability_pred = self.actor(state)

                # Convert stability_pred to match the binary targets if not already
                stability_pred = torch.sigmoid(stability_pred).unsqueeze(1)
                
                # Getting the state value prediction from the critic using the flattened state
                state_value = self.critic(flattened_state).unsqueeze(0)  # Direct application with adjusted state
    
                dist = torch.distributions.Normal(action_means, action_stds)
                action_tensor = torch.tensor(action, dtype=torch.float32, device=action_means.device)
                log_prob = dist.log_prob(action_tensor).sum(axis=-1).unsqueeze(0)
                log_probs_batch.append(log_prob)
                
                # Stability loss
                stability_loss = F.binary_cross_entropy_with_logits(stability_pred, stability_target)
                stability_losses_batch.append(stability_loss.unsqueeze(0))
    
                entropy = dist.entropy().sum(axis=-1).unsqueeze(0)
                entropies_batch.append(entropy)
    
                state_values_batch.append(state_value)  # No additional unsqueeze(0) needed here after correction
            
            # Convert to tensors
            log_probs = torch.cat(log_probs_batch).to(prev_log_probs.device)
            state_values = torch.cat(state_values_batch).squeeze(-1).to(returns.device)  # Ensuring the correct final shape for state_values
            stability_losses = torch.cat(stability_losses_batch)
            entropy = torch.cat(entropies_batch)
    
            # Surrogate function
            ratios = torch.exp(log_probs - prev_log_probs)            
            #advantages_batch = returns - state_values.detach()
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1.0 - self.policy_clip, 1.0 + self.policy_clip) * advantages
            
            actor_loss = - torch.min(surr1, surr2).mean() - self.entropy_coef * entropy.mean() + stability_losses.mean()
            critic_loss = F.mse_loss(state_values, returns)            
            
            # Update actor
            self.optimizer_actor.zero_grad()
            actor_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), max_norm=1.0)
            self.optimizer_actor.step()
            
            # Update critic
            self.optimizer_critic.zero_grad()
            critic_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), max_norm=1.0)
            self.optimizer_critic.step()
            
            # Append the item loss to the lists
            actor_losses.append(actor_loss.item())  # Use item() to extract scalar value
            critic_losses.append(critic_loss.item())  # Use item() to extract scalar value            
   
        return actor_losses, critic_losses

def train(env, agent, num_episodes, max_timesteps, batch_size):    
    for episode in range(start_episode, num_episodes):        
        node_features_tensor, edge_feature_tensor, edge_index, performance_metrics = env.reset()
        state = (node_features_tensor, edge_feature_tensor, edge_index)
        values = []
    
        for t in range(max_timesteps):
            # Selecting action based on the current state
            action, log_prob, perf_metrics = agent.select_action(state, performance_metrics)
            next_node_features_tensor, next_edge_feature_tensor, next_edge_index, reward, done, previous_metrics = env.step(action, perf_metrics)
            next_state = (next_node_features_tensor, next_edge_feature_tensor, next_edge_index)
    
            values.append(agent.critic(state[0]).item())  # Ensure Critic can handle state format
            state = next_state
   
        next_value = agent.critic(state[0]).item() if not done else 0