Skip to content

GitLab

Commit 7cc052f6 authored by Ofey Chan's avatar Ofey Chan Committed by Frank Lee
Browse files

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

parent 46931e3c
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
colossalai/nn/layer/colossalai_layer/linear.py
View file @ 7cc052f6
...@@ -81,7 +81,8 @@ class Linear(ColossalaiModule): ...@@ -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,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment