def METHOD_NAME(device_type, count):
    """Resets logical devices for CPU/GPU.
    Logical devices can only be instantiated once on a particular context. For
    now, context re-use is triggering some function duplication errors, so we
    reset the context on each call.
    Args:
      device_type: The device_type to reset.
      count: numbers of virtual device to reset to.
    """
    if device_type.upper() not in ["CPU", "GPU"]:
        raise ValueError(
            "resetting logical device for non-supported device type: "
            f"{device_type}"
        )
    reset_context()
    cpus = tf.config.list_physical_devices("CPU")
    if device_type.upper() == "GPU":
        gpus = tf.config.list_physical_devices(device_type)
        tf.config.set_logical_device_configuration(
            gpus[0],
            [
                tf.config.LogicalDeviceConfiguration(
                    memory_limit=_DEFAULT_GPU_MEMORY_LIMIT
                ),
            ]
            * count,
        )
    # Always config CPU mesh as the host mesh for DTensor
    tf.config.set_logical_device_configuration(
        cpus[0],
        [
            tf.config.LogicalDeviceConfiguration(),
        ]
        * count,
    )