import torch


def unmerge_second_and_third_dim(
    batch: torch.Tensor, second_dim: int = -1, third_dim: int = -1
) -> torch.Tensor:
    """Method to modify the shape of a batch by unmerging the second and the third dimension.
    Given a tensor of shape (a, b*c, ...), return a tensor of shape (a, b, c, ...)"""
    shape = batch.shape
    return batch.view(second_dim, third_dim, *shape[1:])