NEMO

class pretrain.models.NEMO(wvf_base_channels, wvf_dropout, wvf_dim_rep, acg_base_channels, acg_dropout, acg_dim_rep, dim_embed)

Bases: torch.nn.modules.module.Module

NEMO bimodal contrastive learning model for IBL brain region pretraining [Yu et al., 2025].

Reference implementation: Haansololfp/NEMO.

Interfaces:

forward(wf, acg) -> (z_wf, z_acg) L2-normalised projections representation(wf, acg) -> (wf_rep, acg_rep) pre-projection

representation(wf, acg)

Return pre-projection representations (used for linear probe).

Return type:

tuple[Tensor, Tensor]

forward(wf, acg)

Return L2-normalised post-projection embeddings.

Return type:

tuple[Tensor, Tensor]