Checkpointable: Fix CPU/GPU device placement issues
authorAllen Lavoie <allenl@google.com>
Sat, 3 Mar 2018 00:06:24 +0000 (16:06 -0800)
committerTensorFlower Gardener <gardener@tensorflow.org>
Sat, 3 Mar 2018 00:16:43 +0000 (16:16 -0800)
commit284dac189dcae46c77f1ec70055b13e69c31e4c0
tree51cb5a4e940222cbb9e8a45e4f80ffd42a955326
parentc12f0c5f84699835f9b8111299febf9fc7aba343
Checkpointable: Fix CPU/GPU device placement issues

Restore ops go on the CPU, then the value gets copied to whichever device it needs to be on. This I need to do manually for restores passed as initial_values; for regular save/restore it's done by the SaveableObjects for variables.

Also explicitly places some counters on the CPU.

Adds a GPU-using test for Checkpointable usage.

PiperOrigin-RevId: 187683050
tensorflow/contrib/eager/python/BUILD
tensorflow/contrib/eager/python/checkpointable_utils.py
tensorflow/contrib/eager/python/checkpointable_utils_test.py
tensorflow/python/BUILD
tensorflow/python/training/checkpointable.py