def get_tf_dlrm_v1_config()

in dags/solutions_team/configs/tensorflow/solutionsteam_tf_nightly_supported_config.py [0:0]


def get_tf_dlrm_v1_config(
    tpu_version: TpuVersion,
    tpu_cores: int,
    tpu_zone: str,
    time_out_in_min: int,
    bottom_mlp: List[int],
    embedding_dim: int,
    train_steps: int,
    extraFlags: str = "",
    project_name: str = Project.CLOUD_ML_AUTO_SOLUTIONS.value,
    runtime_version: str = RuntimeVersion.TPU_VM_TF_NIGHTLY.value,
    is_pod: bool = False,
    is_pjrt: bool = True,
    criteo_dir: str = gcs_bucket.CRITEO_DIR,
    network: str = "default",
    subnetwork: str = "default",
    global_batch_size=16384,