def METHOD_NAME(self, node_feat_dim, update_net_hiddens, node_update_type):
    if node_update_type == NodeUpdateType.GRU:
        self.update_net = nn.GRU(node_feat_dim * 3, node_feat_dim)
    else:
        layer = []
        layer.append(nn.Linear(node_feat_dim * 4, update_net_hiddens[0]))
        for i in range(1, len(update_net_hiddens)):
            layer.append(nn.ReLU())
            layer.append(nn.Linear(update_net_hiddens[i - 1], update_net_hiddens[i]))
        layer.append(nn.ReLU())
        layer.append(nn.Linear(update_net_hiddens[-1], node_feat_dim))
        self.update_net = nn.Sequential(*layer)