github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-25 03:31:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

[tensor]fix test_linear (#826)

Browse Source
This commit is contained in:
Ziyue Jiang
2022-04-21 17:18:56 +08:00
committed by GitHub
parent 1a9e2c2dff
commit 8e6fdb4f29
2 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/tensor/_ops/linear.py
View File

@@ -19,8 +19,9 @@ def colo_linear(types, args, kwargs, pg):
bias = None
else:
bias = kwargs.get('bias', None)
if isinstance(bias, ColoTensor):
bias = bias.torch_tensor()
if isinstance(bias, ColoTensor):
bias = bias.torch_tensor()
# Add communication logic before and after linear call.
if isinstance(weight, ColoTensor):