Pretrain 时 采用GPU预训练出错，采用CPU可以

[wanghao0225]

提问时请尽可能提供如下信息：

基本信息

• 你使用的操作系统: win 10
• 你使用的Python版本: 3.7.9
• 你使用的Tensorflow版本: 2.2.0
• 你使用的Keras版本: 2.3.1
• 你使用的bert4keras版本: 0.9.5
• 你使用纯keras还是tf.keras: tf.keras
• 你加载的预训练模型:Albert

核心代码

# 请在此处贴上你的核心代码。
# 请尽量只保留关键部分，不要无脑贴全部代码。

os.environ['TF_KERAS'] = '1'
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ['KERAS_BACKEND'] = 'tensorflow'

AdamW = extend_with_weight_decay(Adam, name='AdamW')
AdamWG = extend_with_gradient_accumulation(AdamW, name='AdamWG')
optimizer = AdamWG(
    learning_rate=5e-6,
    weight_decay_rate=0.01,
    exclude_from_weight_decay=['Norm', 'bias'],
    grad_accum_steps=16,
)
train_model.compile(optimizer=optimizer)

输出信息

# 请在此处贴上你的调试输出
错误信息：
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation model_1/Embedding-Position/ReadVariableOp: Could not satisfy explicit device specification '' because the node {{colocation_node model_1/Embedding-Position/ReadVariableOp}} was colocated with a group of nodes that required incompatible device '/job:localhost/replica:0/task:0/device:GPU:0'. All available devices [/job:localhost/replica:0/task:0/device:CPU:0, /job:localhost/replica:0/task:0/device:XLA_CPU:0, /job:localhost/replica:0/task:0/device:GPU:0, /job:localhost/replica:0/task:0/device:XLA_GPU:0]. 
苏神你好，用了你的pretrain方法，用CPU能调试通过，但切到GPU，就会报错，能帮我分析以下原因吗？

自我尝试

不管什么问题，请先尝试自行解决，“万般努力”之下仍然无法解决再来提问。此处请贴上你的努力过程。

[wanghao0225]

#其中采用AdamW 方法，用GPU可以测试通过，但是加上AdamWG 就会报上面的错。

[bojone]

tf 2.2我也不知道，可能要disable掉垃圾的eager模式吧。

[Lufffya]

#其中采用AdamW 方法，用GPU可以测试通过，但是加上AdamWG 就会报上面的错。

大佬解决了问题吗，我在苏神的 CLUE-bert4keras 下的 cmrc2018 任务也遇到了这个问题