special_fn.py 文件源码

python
阅读 28 收藏 0 点赞 0 评论 0

项目:tefla 作者: openAGI 项目源码 文件源码
def fn_device_dependency(name, device=""):
    """Add control deps for name and device."""
    key = name + "_" + device
    outs = []

    def body():
        with tf.control_dependencies(fn_device_dependency_dict()[key]):
            yield outs
            assert outs

            deps = outs
            if isinstance(outs[0], list) or isinstance(outs[0], tuple):
                assert len(outs) == 1
                deps = outs[0]
            fn_device_dependency_dict()[key] = deps

    if device:
        with tf.device(device):
            return body()
    else:
        return body()
评论列表
文章目录


问题


面经


文章

微信
公众号

扫码关注公众号