您如何指定应将特定的Keras层分配给CPU？

Question

我正在研究序列到序列模型，主要遵循教程here。

我模型中的一层太大了，无法容纳GPU内存，但是我仍然希望能够将GPU用于其他层。是否可以仅在Keras中移动该层？

我尝试这样做：

with tf.device('/device:CPU:0'):
  decoder_outputs = Dense(num_decoder_tokens, activation='softmax')(x)

，但似乎不起作用。我仍然用光了内存。

这似乎可行：

with tf.device('/device:CPU:0'):
  model = Model([encoder_inputs, decoder_inputs], decoder_outputs_dense)
  model.compile(optimizer='adam', loss='categorical_crossentropy')

但是它将整个模型移到CPU上，并且训练时间跳了几千小时，所以我不能这样做。

我想Tensorflow / pyTorch中有可能，但是我想知道使用Keras时是否有办法做到这一点。

编辑： 还有另一个question谈论同样的问题，但是那里的公认答案似乎不适用于Keras，而且我没有足够的声誉来评论答案：（

EDIT2： 上面的link 确实可以解决问题。最初看来，它似乎不起作用，但这是因为我正在使用Jupyter笔记本，并且可能忘记了在编辑或类似操作后执行单元的操作。