Skip to content

Commit 9152d5f

Browse files
committed
Merge: [nnUNet/TF2] Update AMP API
2 parents 0107cf2 + d21ef30 commit 9152d5f

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

  • TensorFlow2/Segmentation/nnUNet/runtime

TensorFlow2/Segmentation/nnUNet/runtime/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ def set_tf_flags(args):
5858
tf.config.threading.set_inter_op_parallelism_threads(max(2, (multiprocessing.cpu_count() // hvd.size()) - 2))
5959

6060
if args.amp:
61-
tf.keras.mixed_precision.experimental.set_policy("mixed_float16")
61+
tf.keras.mixed_precision.set_global_policy("mixed_float16")
6262

6363

6464
def is_main_process():

0 commit comments

Comments
 (0)