if ((tensorflow::IsGoogleCudaEnabled() || tensorflow::IsBuiltWithROCm()) &&
      num_virtual_gpus > 0) {
    tensorflow::GPUOptions* gpu_options =
        server_def.mutable_default_session_config()->mutable_gpu_options();
    auto virtual_devices =
        gpu_options->mutable_experimental()->add_virtual_devices();
    for (int i = 0; i < num_virtual_gpus; ++i) {
      virtual_devices->add_memory_limit_mb(200);
    }
  }

    flags->tf_xla_cpu_global_jit = true;
    flags->tf_xla_min_cluster_size = 1;
  } else {
    optimizer_options->set_global_jit_level(tensorflow::OptimizerOptions::OFF);
  }

  auto* gpu_options = config.mutable_gpu_options();
  gpu_options->set_allow_growth(gpu_memory_allow_growth);

  (*config.mutable_device_count())["CPU"] = num_cpu_devices;

  // TODO(b/113217601): This is needed for EagerContext::runner_ to use a