mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-02 01:28:31 +00:00
[test] modify model supporting part of low_level_zero plugin (including correspoding docs)
This commit is contained in:
@@ -53,16 +53,6 @@ def output_transform_fn(x):
|
||||
return dict(output=x)
|
||||
|
||||
|
||||
def output_transform_fn(x):
|
||||
if isinstance(x, KeyedTensor):
|
||||
output = dict()
|
||||
for key in x.keys():
|
||||
output[key] = x[key]
|
||||
return output
|
||||
else:
|
||||
return dict(output=x)
|
||||
|
||||
|
||||
def get_ebc():
|
||||
# EmbeddingBagCollection
|
||||
eb1_config = EmbeddingBagConfig(name="t1", embedding_dim=SHAPE, num_embeddings=SHAPE, feature_names=["f1"])
|
||||
|
Reference in New Issue
Block a user