def toy_train():
    import matplotlib.pyplot as plt

    C = 1
    F = 64
    T = 7
    shift = 2
    ref = torch.randn(1, C, F, T)
    echo = torch.roll(ref, shifts=shift, dims=-1)
    echo[:, :, :, :shift] = 0
    mic = echo

    layer = Attention_Align_Block(channels=C, fdim=F, pdim=64, max_delay_blocks=8)

    optimizer = torch.optim.Adam(layer.parameters(), lr=0.01)
    loss_fn = torch.nn.MSELoss()
    for i in range(1000):
        pred,_ = layer(mic, ref)
        loss = loss_fn(pred, echo)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if i %100 ==0:
            print('loss: ', loss.item())

    with torch.no_grad():
        out, corr = layer(mic, ref)
        plt.imshow(corr[0].detach().cpu().numpy().T)
        plt.show()