如何重写张量流图以将CPU用于所有操作

我已经在多GPU和CPU设置上训练了网络，并将结果模型保存为tensorflow SavedModel。然后，我还有另一个脚本，可以加载生成的模型并运行所需的操作以进行预测，即对模型进行推断。这可以在训练模型的相同设置上进行。

但是，我需要将模型部署为在具有1个CPU且没有GPU的设备上运行。当我尝试运行相同的脚本时，出现以下错误：

InvalidArgumentError（请参阅上面的回溯）：从检查点还原失败。这很可能是由于当前图形与来自检查点的图形之间的不匹配。请确保您没有更改基于检查点的预期图形。原始错误：

无法为操作default_policy_1 / tower_1 / variable分配设备：已将节点default_policy_1 / tower_1 / variable（在restore.py:56处定义）明确分配给/ device：GPU：0，但可用设备为[/ job：localhost / replica：0 / task：0 / device：CPU：0，/ job：localhost /副本：0 / task：0 / device：XLA_CPU：0]。确保设备规格引用的是有效设备。所请求的设备似乎是GPU，但未启用CUDA。[[node default_policy_1 / tower_1 / Variable（在restore.py:56处定义）]]

这看起来很有希望，但是代码根本没有改变我的图，删除了0个节点-删除了要在CPU上运行的操作图tensorflow

通常，删除所有不在CPU上运行的操作似乎并不明智。

我尝试将所有内容包装在一个with tf.device('CPU:0')块中，以及使用config = tf.ConfigProto(device_count={'GPU': 0})但都没有改变错误。

相关代码：

from tensorflow.python.saved_model import loader

input_tensor_key_feed_dict = {'observations': np.array([[23]]), 'prev_action': np.array([0]),
                              'prev_reward': np.array([0]), 'is_training': False}

config = tf.ConfigProto(device_count={'GPU': 0})
with tf.device('CPU:0'):
    with session.Session(None, graph=ops_lib.Graph(), config=config) as sess:
        
        loader.load(sess, tag_set.split(','), saved_model_dir) #error occurs here
        
        outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
        for i, output in enumerate(outputs):
            output_tensor_key = output_tensor_keys_sorted[i]
            print('Result for output key %s:\t%s' % (output_tensor_key, output))