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

[NFC] polish colossalai/nn/layer/colossalai_layer/linear.py (#1556)

pull/1550/head
Ofey Chan 2022-09-08 15:22:10 +08:00 committed by Frank Lee
parent 46931e3c32
commit 7cc052f6c0
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/nn/layer/colossalai_layer/linear.py
View File

@ -81,7 +81,8 @@ class Linear(ColossalaiModule):
else: else:
linear_cls = _parallel_linear[tensor_parallel] linear_cls = _parallel_linear[tensor_parallel]
gather_output = kwargs.pop('gather_output', None) gather_output = kwargs.pop('gather_output', None)
if 'gather_output' in inspect.signature(linear_cls.__init__).parameters.keys(): # gather_out arg is available if 'gather_output' in inspect.signature(
linear_cls.__init__).parameters.keys(): # gather_out arg is available
kwargs['gather_output'] = gather_output kwargs['gather_output'] = gather_output
layer = linear_cls( layer = linear_cls(
in_features, in_features,