@contextlib.contextmanager
def profiler(round_num):
    if ((rounds_per_profile > 0) and ((round_num % rounds_per_profile) == 0)):
        with tf.profiler.experimental.Profile(summary_logdir):
            (yield)
    else:
        (yield)
@contextlib.contextmanager
def profiler(round_num):
    if ((rounds_per_profile > 0) and ((round_num % rounds_per_profile) == 0)):
        with tf.profiler.experimental.Profile(summary_logdir):
            (yield)
    else:
        (yield)