I need to know that, in my previous code, i got the following error

  File d:\opamp circuits\rl_gnn\rl_gnn_cktenv_new2.py:324 in __init__
    self.actor = actor_class(gnn_model.conv2.out_channels, action_dim, std)  # Initialize actor

TypeError: Actor.forward() takes 2 positional arguments but 4 were given

My previous code:
class PPOAgent:
    def __init__(self, actor_class, critic_class, gnn_model, action_dim, bounds_low, bounds_high,
                 lr_actor=3e-4, lr_critic=1e-3, gamma=0.99, lamda=0.95, epsilon=0.2, std=0.0):
        self.actor = actor_class(gnn_model.conv2.out_channels, action_dim, std)  # Initialize actor        
        self.critic = critic_class(gnn_model.conv2.out_channels)  # Initialize critic
        self.gnn_model = gnn_model  # GNN model instance
        self.optimizer_actor = optim.Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = optim.Adam(self.critic.parameters(), lr=lr_critic)
        self.gamma = gamma
        self.lamda = lamda
        self.epsilon = epsilon
        #self.bounds_low = torch.tensor(bounds_low).float()
        #self.bounds_high = torch.tensor(bounds_high).float()
        self.bounds_low = bounds_low
        self.bounds_high = bounds_high
        self.std = std
        self.epochs = 10  # Define the number of epochs for policy update

    def select_action(self, state, edge_index, edge_attr):
        # Pass state through the GNN model first to get state’s embedding
        state_embedding = self.gnn_model(state, edge_index, edge_attr)
        
        # Then, pass the state embedding through the actor network to get action mean and std
        mean, std = self.actor(state_embedding)

        # Rearranging based on the specifications
        rearranged_mean = rearrange_action_output(mean)

        # Scale mean based on action bounds defined
        mean = self.bounds_low + (torch.sigmoid(rearranged_mean) * (self.bounds_high - self.bounds_low))
        dist = Normal(mean, std)
        
        # Sample an action from the distribution and calculate its log probability
        action = dist.sample()

        action = torch.clamp(action, self.bounds_low, self.bounds_high)  # Ensure action is within bounds
        action_log_prob = dist.log_prob(action)
        
        return action.detach(), action_log_prob.detach()

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)
gnn_model = EnhancedGNNModelWithSharedParams(num_node_features, num_edge_features, num_out_features)
agent = PPOAgent(actor, critic, gnn_model, action_dim, bounds_low, bounds_high, lr_actor, lr_critic, gamma, lambda_, epsilon, std)

# Train agent
train(env, agent, env.num_episodes, env.max_timesteps, env.batch_size, env.epsilon)

I made the following upgrade in my previous code,
class PPOAgent:
    def init(self, actor, critic, gnn_model, action_dim, bounds_low, bounds_high,
                 lr_actor=3e-4, lr_critic=1e-3, gamma=0.99, lamda=0.95, epsilon=0.2, std=0.0):
        self.actor = actor  # Actor instance
        self.critic = critic  # Critic instance
        self.gnn_model = gnn_model  # GNN model instance
        self.optimizer_actor = Adam(self.actor.parameters(), lr=lr_actor)
        self.optimizer_critic = Adam(self.critic.parameters(), lr=lr_critic)
        self.gamma = gamma
        self.lamda = lamda
        self.epsilon = epsilon
        self.bounds_low = bounds_low
        self.bounds_high = bounds_high
        self.std = std
        self.epochs = 10

    def select_action(self, state, edge_index, edge_attr):
        state_embedding = self.gnn_model(state, edge_index, edge_attr)
        mean, std = self.actor(state_embedding)
        
        # Assuming rearrange_action_output is a function you have defined elsewhere
        # rearranged_mean = rearrange_action_output(mean)
        # Treat ‘rearranged_mean’ as ‘mean’ if you do not need to rearrange
        
        dist = Normal(mean, std)
        action = dist.sample()
        action = torch.clamp(action, self.bounds_low, self.bounds_high)
        action_log_prob = dist.log_prob(action)
        
        return action.detach(), action_log_prob.detach()

# Create the environment
env = CircuitEnvironment(server_address, username, password, bounds_low, bounds_high, target_metrics, netlist_content)
gnn_model = EnhancedGNNModelWithSharedParams(num_node_features, num_edge_features, num_out_features)
agent = PPOAgent(actor, critic, gnn_model, action_dim, bounds_low, bounds_high, lr_actor, lr_critic, gamma, lambda_, epsilon, std)

# Train agent
train(env, agent, env.num_episodes, env.max_timesteps, env.batch_size, env.epsilon)

I need to know that, do i need any upgrade in the update policy method with in the class PPOAgent to call the actor function with the gnn_model.
class PPOAgent:
    def init(self, actor, critic, gnn_model, action_dim, bounds_low, bounds_high,
                 lr_actor=3e-4, lr_critic=1e-3, gamma=0.99, lamda=0.95, epsilon=0.2, std=0.0):

    def select_action(self, state, edge_index, edge_attr):       
        return action.detach(), action_log_prob.detach()
    
    def compute_gae(self, next_value, rewards, dones, values, gamma=0.99, lambda_=0.95):
        return returns
    
    def update_policy(self, states, actions, log_probs, returns, advantages):
        for epoch in range(self.epochs):
            sampler = BatchSampler(SubsetRandomSampler(range(len(states))), batch_size=64, drop_last=True)
            for indices in sampler:
                sampled_states = torch.stack([states[i] for i in indices])
                sampled_actions = torch.stack([actions[i] for i in indices])
                sampled_log_probs = torch.stack([log_probs[i] for i in indices])
                sampled_returns = torch.stack([returns[i] for i in indices])
                sampled_advantages = torch.stack([advantages[i] for i in indices])

                # Assuming the actor model returns a distribution from which log probabilities can be computed
                mean, new_std = self.actor(sampled_states)
                dist = Normal(mean, new_std.exp())
                new_log_probs = dist.log_prob(sampled_actions)
                
                ratio = (new_log_probs - sampled_log_probs).exp()
                
                # Incorporating stability loss
                stability_losses = []
                for state in sampled_states:
                    stability_loss_val = stability_loss(state, target_stability=1.0)  # Assuming you want all nodes to be stable
                    stability_losses.append(stability_loss_val)
                mean_stability_loss = torch.mean(torch.stack(stability_losses))

                surr1 = ratio * sampled_advantages
                surr2 = torch.clamp(ratio, 1.0 - self.epsilon, 1.0 + self.epsilon) * sampled_advantages
                actor_loss = -torch.min(surr1, surr2).mean() + mean_stability_loss
                critic_loss = F.mse_loss(sampled_returns, self.critic(sampled_states))

                self.optimizer_actor.zero_grad()
                actor_loss.backward()
                self.optimizer_actor.step()

                self.optimizer_critic.zero_grad()
                critic_loss.backward()
                self.optimizer_critic.step()