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

hotfix colotensor _scan_for_pg_from_args (#1276)

pull/1277/head
ver217 2022-07-12 20:46:31 +08:00 committed by GitHub
parent 0cf8e8e91c
commit 7aadcbd070
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/tensor/colo_tensor.py
View File

@ -40,7 +40,7 @@ def _scan_for_pg_from_args(args, kwargs) -> ProcessGroup:
pg = _scan_for_pg_from_args(elem, {}) pg = _scan_for_pg_from_args(elem, {})
if pg is not None: if pg is not None:
return pg return pg
for k, v in kwargs: for k, v in kwargs.items():
if isinstance(v, ColoTensor): if isinstance(v, ColoTensor):
pg = v.get_process_group() pg = v.get_process_group()
return pg return pg