I needed to support display suite, so I had to change the check on build mode to this

  if ($node->build_mode === NODE_BUILD_NORMAL || !is_numeric($node->build_mode)) {
  if ($element['#node']->build_mode === NODE_BUILD_NORMAL || !is_numeric($element['#node']->build_mode)) {