def set_device(self, device_name):
    """
    Set the device before the next fprop to create a new graph on the
    specified device.
    """
    device_name = unify_device_name(device_name)
    self.device_name = device_name
    for layer in self.layers:
      layer.device_name = device_name