Commit 284dac18 authored by Allen Lavoie's avatar Allen Lavoie Committed by TensorFlower Gardener
Browse files

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
parent c12f0c5f
Show whitespace changes
Inline Side-by-side
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment