in benchmarks/horovod-resnet/train_imagenet_resnet_hvd.py [0:0]
def stage(tensors):
"""Stages the given tensors in a StagingArea for asynchronous put/get.
"""
stage_area = data_flow_ops.StagingArea(
dtypes=[tensor.dtype for tensor in tensors],
shapes=[tensor.get_shape() for tensor in tensors],
)
put_op = stage_area.put(tensors)
get_tensors = stage_area.get()
tf.add_to_collection("STAGING_AREA_PUTS", put_op)
return put_op, get_tensors