"vllm/vscode:/vscode.git/clone" did not exist on "bc8a8ce5ec374dd18e86f59be7cb0057a4b21992"
Unverified Commit 908733ac authored by 22quinn's avatar 22quinn Committed by GitHub
Browse files

[Model] Use sigmoid for single-label classification (#18313)


Signed-off-by: default avatar22quinn <33176974+22quinn@users.noreply.github.com>
parent 1a8f68bb
......@@ -242,9 +242,16 @@ class PoolerHead(nn.Module):
if self.softmax:
if isinstance(pooled_data, list):
pooled_data = [F.softmax(data, dim=-1) for data in pooled_data]
pooled_data = [
F.softmax(data, dim=-1)
if data.shape[-1] >= 2 else F.sigmoid(data)
for data in pooled_data
]
else:
if pooled_data.shape[-1] >= 2:
pooled_data = F.softmax(pooled_data, dim=-1)
else:
pooled_data = F.sigmoid(pooled_data)
return pooled_data
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment