code: def device_function(self, var):
      
      if var.type not in ('Variable', 'VariableV2', 'VarHandleOp'):
        tf.logging.debug('Place {} on last device: {}.'.format(
            var.name, self._last_device))
        return self._last_device
  
      shape = tf.TensorShape(var.get_attr('shape'))
      assert shape.num_elements() is not None
  
      size = var.get_attr('dtype').size
      mem, device = heapq.heappop(self._mem_device_heap)
      mem += shape.num_elements() * size
      heapq.heappush(self._mem_device_heap, (mem, device))
      tf.logging.debug('Place variable {} on {} and consumes {} Bytes.'.format(
          var.name, device, mem))
      self._last_device = device
  
      return device
docs: Choose a device for the input variable.
  
      Args:
        var: an Variable.
  
      Returns:
        The device for placing the var.