github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-16 23:16:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/nn/_ops/embedding.py code style (#1561)

Browse Source
This commit is contained in:
BigOneLiXiaoMing 2022-09-08 16:02:44 +08:00 committed by Frank Lee
parent 08815f0e72
commit 0c4c9aa6e0
1 changed files with 10 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/nn/_ops/embedding.py
View File

@ -113,8 +113,7 @@ def colo_embedding(input_tensor: GeneralTensor,
if not weight.has_compute_spec(): # No Model Parallel Applied if not weight.has_compute_spec(): # No Model Parallel Applied
assert weight.is_replicate(), 'Invalid weight spec for native embedding op' assert weight.is_replicate(), 'Invalid weight spec for native embedding op'
return ColoTensor.from_torch_tensor( return ColoTensor.from_torch_tensor(tensor=F.embedding(input_tensor,
tensor=F.embedding(input_tensor,
weight, weight,
padding_idx=padding_idx, padding_idx=padding_idx,
max_norm=max_norm, max_norm=max_norm,