Loading model weights partially

[bilgineroral]

Hello, I really appreciate this amazing work. The problem I face is the models are a bit large to run on limited hardware. My question is, is it possible to partially load the model? For example, I need the embeddings extracted from sequence and coordinates tracks for a downstream task and other tracks are not really even used, so they’re loaded into GPU for no reason. Is it possible to load weights only for specified tracks, which I believe would significantly reduce unnecessary VRAM usage? Thank you so much.

[ebetica]

I believe the encoders and decoders won't be loaded if you never call the encode and decode functions.

esm/esm/models/esm3.py

Lines 222 to 224 in 30eb907

	self._structure_encoder = None
	self._structure_decoder = None
	self._function_decoder = None

Otherwise, try wrapping the model forwards in a torch.no_grad.