def METHOD_NAME(self, x, messages):
    if self.node_update_type == NodeUpdateType.GRU:
        _, new_x = self.update_net(messages.unsqueeze(0), x.unsqueeze(0))
        new_x = torch.squeeze(new_x)
    else:
        update_inputs = torch.cat((messages, x), dim=-1)
        new_x = self.update_net(update_inputs)
    if hasattr(self, 'update_norm'):
        new_x = self.update_norm(new_x)
    if self.node_update_type == NodeUpdateType.RESIDUAL:
        return x + new_x
    return new_x