mirror of
https://github.com/dmlc/dgl.git
synced 2026-06-04 19:44:23 +08:00
74c9d27d16
* [Misc] Auto-format tests. * more --------- Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
38 lines
978 B
Python
38 lines
978 B
Python
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()
|