'Returns the device to use for caching variables.
Variables are cached on the worker CPU when using replicas.
Returns:
A device string or None if the variables do not need to be cached.'

def caching_device(self):

    if (self._num_ps_tasks > 0):
        return (lambda op: op.device)
    else:
        return None
