def get_device(mdl):
    """
    Checks the device of the first set of params.

    https://discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda/180
    :return:
    """
    device = next(mdl.parameters()).device
    return device