Skip to content
Snippets Groups Projects
Commit 86ac08b5 authored by Alexandre Chapin's avatar Alexandre Chapin :race_car:
Browse files

Fix out proj

parent cdf3ad56
No related branches found
No related tags found
No related merge requests found
...@@ -238,6 +238,7 @@ class Attention(nn.Module): ...@@ -238,6 +238,7 @@ class Attention(nn.Module):
# Get output # Get output
out = attn @ v out = attn @ v
out = self._recombine_heads(out) out = self._recombine_heads(out)
self.out_proj = self.out_proj.to(out.device)
out = self.out_proj(out) out = self.out_proj(out)
return out return out
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment