import torch
import inspect


def _move_prediction_device(predictions, device):  
    """ Move a prediction to specified device. This function may modify predictions in place
    
    Args:
        predictions: the original prediction.
        device (torch Device): the device to move the prediction to.  
        
    Returns:
        prediction: the prediction in the new device
    """
    if issubclass(type(predictions), torch.distributions.distribution.Distribution):  # Trick to set the device of a torch Distribution class because there is no interface for this
        for (name, value) in inspect.getmembers(predictions, lambda v: isinstance(v, torch.Tensor)):
            try:
                setattr(predictions, name, value.to(device))
            except:   # This is a hack, some properties cannot be set because they are calculated from other properties
                pass
        return predictions
    elif issubclass(type(predictions), dict):  # If dictionary then move each element to device separately 
        for key, item in predictions.items():
            predictions[key] = _move_prediction_device(item, device) 
        return predictions
    else:
        return predictions.to(device)