import torch


def get_devices(x):
    """Recursively find all devices in state dictionary."""
    devices = set()
    if isinstance(x, dict):
        for key, val in x.items():
            devices.update(get_devices(val))
    else:
        if isinstance(x, torch.Tensor):
            if x.device.type == 'cpu':
                devices.add(-1)
            else:
                devices.add(x.device.index)
    return devices