nlpstack.integrations.torch.modules.heads module#