import torch


def _get_prediction_device(predictions):
    """ Get the device of a prediction
    
    Args:
        predictions: a prediction of any type. 
        
    Returns:
        device: the torch device that prediction is on. 
    """
    if issubclass(type(predictions), torch.distributions.distribution.Distribution):
        with torch.no_grad():  
            device = predictions.sample().device    # Trick to get the device of a torch Distribution class because there is no interface for this
    elif issubclass(type(predictions), dict):
        assert len(predictions.keys()) != 0, "Must have at least one element in the ensemble"
        device = _get_prediction_device(predictions[next(iter(predictions))])   # Return the device of the first element in the dictionary 
    else:
        device = predictions.device
    return device