def load_data(data, device): 
    """
    Loads onto appropriate device (cpu or gpu)
    Args: 
        data: mini-batch data from dataloader 
        device: <torch.device> 
    Returns: 
        input_data 
        target_data
    """
    if device.type == 'cuda':
        input_data, target_data = data[0].to(device), data[1].to(device)
    else: 
        input_data, target_data = data[0], data[1]
    return input_data, target_data