<p dir="auto"><strong>System information</strong></p>
<ul dir="auto">
<li>Have I written custom code (as opposed to using a stock example script provided in TensorFlow):<br>
Yes, duplicated the minimal code from documentation snippet from <a href="https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/distribute#example-with-keras-api">here</a></li>
<li>OS Platform and Distribution (e.g., Linux Ubuntu 16.04):"18.04.1 LTS (Bionic Beaver)"</li>
<li>Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:None</li>
<li>TensorFlow installed from (source or binary): source and binary</li>
<li>TensorFlow version (use command below): v1.12.0-rc2-0-g748435b8ef</li>
<li>Python version: 3.6.6</li>
<li>Bazel version (if compiling from source): 0.15.2</li>
<li>GCC/Compiler version (if compiling from source): 6.4.0</li>
<li>CUDA/cuDNN version: 9.0 / 7.3</li>
<li>GPU model and memory: 2 x 1080-ti</li>
</ul>
<p dir="auto"><strong>Describe the current behavior</strong><br>
Cannot save a tf.keras model if trained with MirroredStrategy either by calling save_weight or by a tf.keras.callbacks.ModelCheckpoint, but does work if MirroredStrategy is not used.</p>
<p dir="auto"><strong>Code to reproduce the issue</strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)

features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
train_dataset = tf.data.Dataset.zip((features, labels))

distribution = tf.contrib.distribute.MirroredStrategy()

model.compile(loss='categorical_crossentropy',
              optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2),
              distribute=distribution)
model.fit(train_dataset, epochs=5, steps_per_epoch=10)"><pre class="notranslate"><code class="notranslate">inputs = tf.keras.layers.Input(shape=(1,))
predictions = tf.keras.layers.Dense(1)(inputs)
model = tf.keras.models.Model(inputs=inputs, outputs=predictions)

features = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
labels = tf.data.Dataset.from_tensors([1.]).repeat(10000).batch(10)
train_dataset = tf.data.Dataset.zip((features, labels))

distribution = tf.contrib.distribute.MirroredStrategy()

model.compile(loss='categorical_crossentropy',
              optimizer=tf.train.GradientDescentOptimizer(learning_rate=0.2),
              distribute=distribution)
model.fit(train_dataset, epochs=5, steps_per_epoch=10)
</code></pre></div>
<p dir="auto">And adding:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="model.save_weights('my_weight')"><pre class="notranslate"><code class="notranslate">model.save_weights('my_weight')
</code></pre></div>
<p dir="auto">Error looks like:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py&quot;, line 72, in get
    return self._index[device]
KeyError: '/replica:0/task:0/device:CPU:0'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File &quot;minimal_example.py&quot;, line 17, in &lt;module&gt;
    model.save_weights('my_weight')
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py&quot;, line 1449, in save_weights
    session = backend.get_session()
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py&quot;, line 469, in get_session
    _initialize_variables(session)
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py&quot;, line 722, in _initialize_variables
    variables = _get_variables(ops.get_default_graph())
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py&quot;, line 716, in _get_variables
    variables.update(opt.optimizer.variables())
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py&quot;, line 787, in variables
    optimizer_variables = [v for v in self._non_slot_variables()
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py&quot;, line 788, in &lt;listcomp&gt;
    if _from_current_graph(v)]
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py&quot;, line 781, in _from_current_graph
    return variable.op.graph is current_graph
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py&quot;, line 308, in op
    return self.get().op
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py&quot;, line 76, in get
    (device, self._index.keys(), device_util.current())), e)
  File &quot;&lt;string&gt;&quot;, line 3, in raise_from
ValueError: Device /replica:0/task:0/device:CPU:0 not found in dict_keys(['/replica:0/task:0/device:GPU:0', '/replica:0/task:0/device:GPU:1']) (current device )"><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py", line 72, in get
    return self._index[device]
KeyError: '/replica:0/task:0/device:CPU:0'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "minimal_example.py", line 17, in &lt;module&gt;
    model.save_weights('my_weight')
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1449, in save_weights
    session = backend.get_session()
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py", line 469, in get_session
    _initialize_variables(session)
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py", line 722, in _initialize_variables
    variables = _get_variables(ops.get_default_graph())
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/backend.py", line 716, in _get_variables
    variables.update(opt.optimizer.variables())
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py", line 787, in variables
    optimizer_variables = [v for v in self._non_slot_variables()
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py", line 788, in &lt;listcomp&gt;
    if _from_current_graph(v)]
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/training/optimizer.py", line 781, in _from_current_graph
    return variable.op.graph is current_graph
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py", line 308, in op
    return self.get().op
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/contrib/distribute/python/values.py", line 76, in get
    (device, self._index.keys(), device_util.current())), e)
  File "&lt;string&gt;", line 3, in raise_from
ValueError: Device /replica:0/task:0/device:CPU:0 not found in dict_keys(['/replica:0/task:0/device:GPU:0', '/replica:0/task:0/device:GPU:1']) (current device )
</code></pre></div>
<p dir="auto">Changing to:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="checkpoint_path = &quot;my_weight&quot;
cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1,
                                                 period=1)
model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=10,
          callbacks=[cp_callback])"><pre class="notranslate"><code class="notranslate">checkpoint_path = "my_weight"
cp_callback = tf.keras.callbacks.ModelCheckpoint(checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1,
                                                 period=1)
model.fit(train_dataset,
          epochs=5,
          steps_per_epoch=10,
          callbacks=[cp_callback])
</code></pre></div>
<p dir="auto">Also does not work an end up with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Epoch 1/5
 1/10 [==&gt;...........................] - ETA: 4s - loss: 1.1921e-07
Epoch 00001: saving model to model_dir/my_weight
WARNING:tensorflow:You are accessing attribute _replicated_modelof the DistributedCallbackModel that may not have been set correctly.
Traceback (most recent call last):
  File &quot;minimal_example.py&quot;, line 35, in &lt;module&gt;
    callbacks=[cp_callback])
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py&quot;, line 1624, in fit
    validation_steps=validation_steps)
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_distributed.py&quot;, line 198, in fit_loop
    callbacks.on_epoch_end(epoch, epoch_logs)
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/callbacks.py&quot;, line 214, in on_epoch_end
    callback.on_epoch_end(epoch, logs)
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/callbacks.py&quot;, line 599, in on_epoch_end
    self.model.save_weights(filepath, overwrite=True)
  File &quot;/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py&quot;, line 2336, in save_weights
    self._replicated_model.save_weights(filepath, overwrite=overwrite,
AttributeError: 'NoneType' object has no attribute 'save_weights'"><pre class="notranslate"><code class="notranslate">Epoch 1/5
 1/10 [==&gt;...........................] - ETA: 4s - loss: 1.1921e-07
Epoch 00001: saving model to model_dir/my_weight
WARNING:tensorflow:You are accessing attribute _replicated_modelof the DistributedCallbackModel that may not have been set correctly.
Traceback (most recent call last):
  File "minimal_example.py", line 35, in &lt;module&gt;
    callbacks=[cp_callback])
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 1624, in fit
    validation_steps=validation_steps)
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_distributed.py", line 198, in fit_loop
    callbacks.on_epoch_end(epoch, epoch_logs)
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/callbacks.py", line 214, in on_epoch_end
    callback.on_epoch_end(epoch, logs)
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/callbacks.py", line 599, in on_epoch_end
    self.model.save_weights(filepath, overwrite=True)
  File "/home/BP/anaconda3/envs/tf12_gpu/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 2336, in save_weights
    self._replicated_model.save_weights(filepath, overwrite=overwrite,
AttributeError: 'NoneType' object has no attribute 'save_weights'
</code></pre></div>
<hr>
<h3 dir="auto">System information</h3>
<ul dir="auto">
<li><strong>Have I written custom code (as opposed to using a stock example script provided in TensorFlow)</strong>:<br>
No</li>
<li><strong>OS Platform and Distribution (e.g., Linux Ubuntu 16.04)</strong>:<br>
N/A</li>
<li><strong>Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device</strong>:<br>
N/A</li>
<li><strong>TensorFlow installed from (source or binary)</strong>:<br>
N/A</li>
<li><strong>TensorFlow version (use command below)</strong>:<br>
N/A</li>
<li><strong>Python version</strong>:<br>
N/A</li>
<li><strong>Bazel version (if compiling from source)</strong>:<br>
N/A</li>
<li><strong>GCC/Compiler version (if compiling from source)</strong>:<br>
N/A</li>
<li><strong>CUDA/cuDNN version</strong>:<br>
N/A</li>
<li><strong>GPU model and memory</strong>:<br>
N/A</li>
<li><strong>Exact command to reproduce</strong>:<br>
N/A</li>
</ul>
<h3 dir="auto">Describe the problem</h3>
<p dir="auto">We need to expose <code class="notranslate">align_corners</code> in this Keras upsample<br>
<a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/convolutional.py#L1935">https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/convolutional.py#L1935</a></p>
<p dir="auto">See also <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="354016520" data-permission-text="Title is private" data-url="https://github.com/keras-team/keras/issues/10994" data-hovercard-type="pull_request" data-hovercard-url="/keras-team/keras/pull/10994/hovercard?comment_id=430305648&amp;comment_type=issue_comment" href="https://github.com/keras-team/keras/pull/10994#issuecomment-430305648">keras-team/keras#10994 (comment)</a></p>