import torch


def get_prediction(model, batch, device):
    """Get predicted labels for given input batch and model."""
    images = torch.tensor(batch, dtype=torch.float).to(device)
    outputs = model(images)
    _, predicted = torch.max(outputs.data, 1)
    return predicted