import torch


def batched_block_diag_forloop(m1, m2, f):
    """
    Args:
        m1, m2: 3-d tensors

        f: either torch.block_diag or anti_block_diag
    """
    bs, _, _ = m1.size()
    res_list = []
    for i in range(bs):
        res_list.append(f(m1[i], m2[i]))
    res = torch.stack(res_list, dim=0)
    return res