Another potential solution, at least for image tensors
def unstack(tensor: torch.Tensor, dimensions: int):
    """
    Unstacks the provided image tensor into the specified amount of dimensions. 
    Expects a tensor in the shape of bs x c x n x n, where bs is the batch size,
    c is the number of channels. Unpacks the image to the size 
    (n / dimensions) ^2 * bs, c, dimensions, dimensions

    Parameters:
    - tensor (torch.Tensor): Image tensor to unstack
    - dimensions (int): Dimensions to unpack to.

    Returns
    - tensor (torch.Tensor): Unstacked image tensor
    """

    assert tensor.shape[-2] == tensor.shape[-1], "Image tensor must be square" 
    assert tensor.shape[-1] % dimensions == 0, "Cannot unpack the image into that shape"

    # Unfold the tensor to the specified dimensions
    tensor = tensor.unfold(1, 1, 1).unfold(2,dimensions,dimensions).unfold(3,dimensions,dimensions)

    # Reshape the tensor to be the shape
    tensor = tensor.squeeze().permute(2,1,0,3,4).reshape(-1, 3, dimensions,dimensions)

    return tensor