import tensorflow as tf

        return float(
            self.sess.run(self.cross_entropy, feed_dict={self.x: xs, self.y_: ys})
        )

    def grad(self, xs, ys):
        """Computes the gradients of the network."""
        return self.sess.run(
            self.cross_entropy_grads, feed_dict={self.x: xs, self.y_: ys}
        )


@ray.remote
class NetActor(object):
    def __init__(self, xs, ys):
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
        with tf.device("/cpu:0"):
            self.net = LinearModel([784, 10])
            self.xs = xs
            self.ys = ys

    # Compute the loss on a batch of data.
    def loss(self, theta):
        net = self.net
        net.variables.set_flat(theta)
        return net.loss(self.xs, self.ys)

    # Compute the gradient of the loss on a batch of data.
    def grad(self, theta):
        net = self.net
        net.variables.set_flat(theta)
