github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Tensor] apply ColoTensor on Torch functions (#821) 

* Revert "[zero] add ZeroTensorShardStrategy (#793)"

This reverts commit 88759e289e.

* [gemini] set cpu memory capacity

* [log] local throughput collecting

* polish

* polish

* polish

* polish code

* polish

* polish code

* add a new tensor structure and override linear for it

* polish

* polish

* polish

* polish

* polish

* polish

* polish

* polish

* polish

* polish

* polish

* [tensor] renaming and reorganize directory structure.

* rm useless dir

* polish

* polish

* [tensor] hander the function not wrapped
pull/824/head
Jiarui Fang 2022-04-21 14:21:10 +08:00 committed by GitHub
parent 0ce8924ceb
commit 660d2d1f1b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_tensor/test_op.py
View File

@ -54,6 +54,13 @@ def test_element_wise():
assert allclose(torch.nn.functional.relu(t), torch.nn.functional.relu(t_ref))
# Test a function not wrapped by
def test_no_wrap_op():
t_ref = torch.randn(3, 5)
t = ColoTensor(t_ref.clone())
assert torch.sum(t) == torch.sum(t_ref)
if __name__ == '__main__':
test_linear()
test_no_wrap_op()
# test_element_wise()