how to add projection

[threestone965]

How to add a projection module to an existing model? The hidden_state of this model is too large.
i added projection=128 in the train_embonly.sh, but the result files don't not contain the projection weights.

[Muennighoff]

The projection layer is created here

gritlm/gritlm/gritlm.py

Line 43 in a73526e

self.projection = torch.nn.Linear(

and should be created when you pass the projection kwarg

It should get saved when the model is saved i.e. here

gritlm/gritlm/training/run.py

Line 422 in a73526e

trainer.save_model()

but if not you can e.g. add something like torch.save(model.projection, "projection.pt") right there below to save it - Let me know how it goes!