e8b348915f3fe67d79f75c4399606821b413e5a1,sonnet/src/rmsprop_test.py,RMSPropTest,testMovingVariablesColocatedWithOriginal,#RMSPropTest#Any#,113

Before Change

    device = "GPU:0" if self.primary_device == "GPU" else "CPU:0"
    with tf.device(device):
      var = tf.Variable(1.0)
    for moving_var in optimizer._get_or_create_moving_vars(var):
      if moving_var is not None:
        self.assertEqual(moving_var.device, var.device)

if __name__ == "__main__":
  // tf.enable_v2_behavior()
  tf.test.main()

After Change

    optimizer = opt_class(learning_rate=0.1)
    with tf.device("CPU:0"):
      var = tf.Variable(1.0)
    optimizer.apply([tf.constant(0.1)], [var])
    self.assertEqual(optimizer.mom[0].device, var.device)
    self.assertEqual(optimizer.ms[0].device, var.device)

Instances