i need you to check my complete code of my PPO agent and check any of the computation were missed it include in the code, check any computation were implemented, but it doesnot included in the following process/procedure, check any data missing or any variables/computed variable missing and check all the necessary errors in the below given class 'PPOAgent' completely.

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 select_action(self, state, performance_metrics):               
        
        action_means, action_stds = self.actor(state) 
        #print("action_means", action_means)

        dist = torch.distributions.Normal(action_means, action_stds)
        action = dist.sample()        
        
        log_probs = dist.log_prob(action).sum(axis=-1)  # Summing over actions if action space is multi-dimensional
        
        # Using tanh to squash gnn_output to [-1, 1]
        scaled_action = torch.tanh(action)
        
        # Now scale and shift the normalized output to the action bounds
        scaled_actions = self.actor.bounds_low + (self.actor.bounds_high - self.actor.bounds_low) * ((scaled_action + 1) / 2)

        return scaled_actions.detach().numpy(), log_probs.detach(), performance_metrics

    def compute_gae(self, next_value, rewards, masks, values):
        # Generalized Advantage Estimation (GAE)
        returns = []
        gae = 0
        for step in reversed(range(len(rewards))):
            delta = rewards[step] + self.gamma * next_value * masks[step] - values[step]
            gae = delta + self.gamma * self.gae_lambda * masks[step] * gae
            returns.insert(0, gae + values[step])
            next_value = values[step]
        return returns

    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 = [], [], []
    
            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
       
                action_means, action_stds = self.actor(state)
                
                # 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)
 
                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
            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()
            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
            
        # Calculate average loss over all epochs
        average_loss_actor = sum(actor_losses) / len(actor_losses)
        average_loss_critic = sum(critic_losses) / len(critic_losses)
    
        return average_loss_actor, average_loss_critic