def model_device(model):
    """Determine the device the model is allocated on."""
    # Source: https://discuss.pytorch.org/t/how-to-check-if-model-is-on-cuda/180
    if next(model.parameters()).is_cuda:
        return 'cuda'
    return 'cpu'