def context_of(data):
    """Return the device of the data which can be either a tensor or a list/dict of tensors."""
    if isinstance(data, Mapping):
        return F.context(next(iter(data.values())))
    elif is_listlike(data):
        return F.context(next(iter(data)))
    else:
        return F.context(data)
Return the device of the data which can be either a tensor or a list/dict of tensors.
python/dgl/utils/internal.py