<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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