Skip to content

GitLab

Commit 32974988 authored by David Chen's avatar David Chen Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 267486140
parent 4bd15fa6
Show whitespace changes
Inline Side-by-side
Showing with 4 additions and 1 deletion
official/recommendation/ncf_keras_main.py
View file @ 32974988
......@@ -254,7 +254,10 @@ def run_ncf(_):
"val_HR_METRIC", desired_value=FLAGS.hr_threshold)
callbacks.append(early_stopping_callback)
with tf.device(tpu_lib.get_primary_cpu_task(params["use_tpu"])):
use_remote_tpu = params["use_tpu"] and FLAGS.tpu
primary_cpu_task = tpu_lib.get_primary_cpu_task(use_remote_tpu)
with tf.device(primary_cpu_task):
(train_input_dataset, eval_input_dataset,
num_train_steps, num_eval_steps) = \
(ncf_input_pipeline.create_ncf_input_data(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment