Based on Summary Trace API,
device_name = tf.test.gpu_device_name()
if not tf.test.is_gpu_available():
raise SystemError('GPU device not found')
print('Found GPU at: {}'.format(device_name))
os.makedirs(os.path.join(args.exp_dir, 'plugins/profile'), exist_ok=True)
tf.summary.trace_on(graph=True, profiler=True)
tracing_params = params.copy()