The same feature dimension limitation when using PGExplainer

When i use the PGExplainer, i discover a problem. i have to use the same feature dimension in the last gnn layer to fit the readout function and run correctly. The example is follows:
class Model(nn.Module):
def __init__(self, in_feats, out_feats):
super().__init__()
self.conv = GraphConv(in_feats, out_feats)
self.fc = nn.Linear(out_feats, out_feats)
nn.init.xavier_uniform_(self.fc.weight)

def forward(self, g, h, embed=False, edge_weight=None):
h = self.conv(g, h, edge_weight=edge_weight)
if embed:
return h
with g.local_scope():
g.ndata["h"] = h
hg = dgl.mean_nodes(g, "h")
return self.fc(hg)

if the input feature dimension is not equals to the output feature dimension in self.fc. it will error. it is not convenient to design the feature size in last gnn dimension layer and JK-Net or skip connection too. how can i solve the problem?