test_pin_memory.py

import backend as F

import dgl
import pytest


@pytest.mark.skipif(
    F._default_context_str == "cpu", reason="Need gpu for this test"
)
def test_pin_unpin():
    t = F.arange(0, 100, dtype=F.int64, ctx=F.cpu())

    assert not F.is_pinned(t)

    if F.backend_name == "pytorch":
        nd = dgl.utils.pin_memory_inplace(t)
        assert F.is_pinned(t)
        nd.unpin_memory_()
        assert not F.is_pinned(t)
        del nd

        # tensor will be unpinned immediately if the returned ndarray is not saved
        dgl.utils.pin_memory_inplace(t)
        assert not F.is_pinned(t)

        t_pin = t.pin_memory()
        # cannot unpin a tensor that is pinned outside of DGL
        with pytest.raises(dgl.DGLError):
            F.to_dgl_nd(t_pin).unpin_memory_()
    else:
        with pytest.raises(dgl.DGLError):
            # tensorflow and mxnet should throw an error
            dgl.utils.pin_memory_inplace(t)


if __name__ == "__main__":
    test_pin_unpin()