import torch


def init_stack_2d(batch_size, seq_length, stack_depth, stack_width, dvc='cpu'):
    """
    Initialization of the stack state.

    Returns
    -------
    stack: tensor
        tensor filled with zeros
    """
    return torch.zeros(batch_size, seq_length, stack_depth, stack_width).to(dvc)