How to implement edge feature in message passing algorithm of Pytorch Geometric?

Question

I am trying to implement a custom message passing function $ x^k_i = x^{k-1}_i + m^k_i$ where $x_i$ is the node embedding of the i'th node and $m^k_i$ is the message getting delivered at the i'th node from its neighbours.

The message is defined as $m^k_i = \frac{1}{N} \sum_{j\in N(i)} ReLU(z_{ij})$ , where $ z_{ij} = x^{k-1}i + x^{k-1}j + \sigma(w*e_{ij} + b)$ ?

Where $x_j$ = neighbouring node and $e_{ij} = $ Edge feature between node 'i' and 'j'.

Code:

import torch
from torch.nn import Sequential as Seq, Linear, ReLU
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import remove_self_loops, add_self_loops

class GC_a(MessagePassing):
    def _init_(self, in_channels, out_channels, hidden_channel, edge_feature):
        super(GC_a, self)._init_(aggr='mean') 
        self.lin_e = torch.nn.Linear(edge_feature, hidden_channel)
        self.act_e = torch.nn.ReLU()

        
    def forward(self, x, edge_index, edge_feature):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]
        edge_index, _ = remove_self_loops(edge_index)
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        
        return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x)

    def message(self, x_i, x_j, e_ij): # e_ij is edge feature ()


        z_ij = torch.cat([x_i, x_j], dim=1) # These are the first 2 terms of the message expression; z_ij has shape [in_channels+in_channels,1] 
        
        e_f = self.lin_e(e_ij) # e_ij [edge_feature,1]
        e_f = self.act_e(e_f) # [hidden_channel,1]

        z_ij = torch.cat([z_ij, e_f], dim=1) # z_ij has shape [in_channels+in_channels+hidden_channels,1]

        return z_ij

    def update(self, aggr_out, x_i):
        
        # aggr_out has shape [N, out_channels]
        new_embedding = torch.cat([aggr_out], dim=1)
        
        return new_embedding