import torch


def batch_rotate_p4(batch, k, device=None):
    """Rotates by k*90 degrees each sample in a batch.
    Args:
        batch (Tensor): the batch to rotate, format is (N, C, H, W).
        k (list of int): the rotations to perform for each sample k[i]*90 degrees.
        device (str): the device to allocate memory and run computations on.
    
    Returns (Tensor):
        The rotated batch.
    """
    batch_size = batch.shape[0]
    assert len(k) == batch_size, "The size of k must be equal to the batch size."
    batch_p4 = torch.empty_like(batch).to(device)
    for i in range(batch_size):
        batch_p4[i] = torch.rot90(batch[i], k=int(k[i]), dims=(1,2))
    return batch_p4