1f0750670cf8ea24ad264debf9603002ab0fb565,test/nn/pool/test_mem_pool.py,,test_mem_pool,#,5

Before Change

    mpool = MemPool(heads=3, num_keys=2, in_channels=2, out_channels=3)
    assert mpool.__repr__() == "MemPool(2, 3, heads=3, num_keys=2)"
    x = torch.rand((5, 4, 2))
    mask = torch.Tensor([[1, 1, 1, 0], [1, 1, 1, 1], [1, 1, 1, 1],
                         [1, 1, 1, 1], [1, 1, 1, 0]]).bool()
    x, S = mpool(x, mask)
    loss, P = MemPool.kl_loss(S, mask)
    assert x.shape == torch.Size([5, 2, 3])
    assert (S[~mask] == 0).all()
    assert P.shape == S.shape

After Change

    assert out.size() == (5, 2, 8)
    assert S[~mask].sum() == 0
    assert S[mask].sum() == x.size(0)
    assert float(loss) > 0

Instances