I am getting the following error, please debug the error in the below code

values.append(agent.critic(node_features_tensor, edge_index).item())  # Ensure Critic can handle state format

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x480 and 24x128)

Description of my process:
In my graph network(undirected bipartite), i have fixed topology of component nodes(11 nodes) on one side and net node(9 nodes) on another side(derived from the analog circuit netlist). where i have total 20 nodes and each nodes are having 24 node features and the 38 edges were connected between these 11 component nodes and 9 net nodes (this is my state representation the data from the environment for optimization process) and among the node features of each node where the 24th node represents its stability state. my environment need the action space of values for 13 variables from the actor network as input to the simulator circuit netlist. here few nodes [0, 3, 5, 7, 10, 13, 15, 16] are need to be check for its stability and during optimization process these nodes are need to be keep in stable state.
Now i modified my critic network to utilize the same gnn_model for finding the value function, which is used by the actor network. so i got this error.



# Initialize Hyperparameters
in_channels = 24  # Number of input features
out_channels = 24  # Number of output actions or features
state_dim = 128  # Dimension of state representation, to match critic input
action_space = 13  # Environment’s action space

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)

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(out_channels, 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.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

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)
    
        episode_rewards = []
        states = []
        actions = []
        log_probs = []
        values = []
        masks = []
        perf_metrics_log = []
    
        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)
    
            episode_rewards.append(reward)
            states.append(state)
            actions.append(action)
            log_probs.append(log_prob.item())
            values.append(agent.critic(node_features_tensor, edge_index).item())  # Ensure Critic can handle state format
            masks.append(1.0 - float(done))
            perf_metrics_log.append(perf_metrics)
                
            performance_metrics = previous_metrics
            state = next_state
    
            if done:
                break