Failed to Get Trivial Embedding

When I try to extract a trivial embedding with the provided function `task_similarity.get_trivial_embedding_from(e)` where e is an embedding, I get an error message `TypeError: 'Embedding' object is not subscriptable`. I notice that in task2vec.py, Embedding is a class and does not have keys such as 'layers', which is why I cannot extract the trivial embedding.

My question is:
1. Is my input of an 'Embedding' object to the function correct?
2. Is there a way to work around this?

Thanks.