dmlc--dgl
nv-dlasalle
1f2e696080
* Disable pinning non-contiguous memory * Prevent views from being converted for write * Fix linting * Add unit tests * Improve error message for users * Switch to pytest function * exclude mxnet and tensorflow from inplace pinning * Add skip * Restrict to pytorch backend * Use backend to retrieve device * Fix capitalization in decorator Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>
23 行
622 B
Python
23 行
622 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':
|
|
dgl.utils.pin_memory_inplace(t)
|
|
assert F.is_pinned(t)
|
|
dgl.utils.unpin_memory_inplace(t)
|
|
assert not F.is_pinned(t)
|
|
else:
|
|
with pytest.raises(dgl.DGLError):
|
|
# tensorflow and mxnet should throw an erro
|
|
dgl.utils.pin_memory_inplace(t)
|
|
|
|
if __name__ == "__main__":
|
|
test_pin_unpin()
|