I am using the pre-trained model of Biomegatron. I want to get word embeddings by dropping the last layer. But I couldn’t figure out how can I do that with Megatron models or Nemo models. I would be appreciated if you help me I’m digging that for 3 days tnx.
Usually you can “drop” layers by replacing them with nn.Identity modules, which would allow you to just reuse the forward method without overriding it.
print(model) should give you all layer names (or alternatively check the model implementation in the corresponding repository).