Update modeling_tiny_mixtral.py
Browse files- modeling_tiny_mixtral.py +1 -0
modeling_tiny_mixtral.py
CHANGED
|
@@ -340,6 +340,7 @@ class SimpleMultiHeadAttention(nn.Module):
|
|
| 340 |
|
| 341 |
# Reshape back to (batch_size, seq_len, dim)
|
| 342 |
y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, self.dim)
|
|
|
|
| 343 |
|
| 344 |
# Output projection
|
| 345 |
y = self.resid_dropout(self.c_proj(y))
|
|
|
|
| 340 |
|
| 341 |
# Reshape back to (batch_size, seq_len, dim)
|
| 342 |
y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, self.dim)
|
| 343 |
+
y = y.to(self.device)
|
| 344 |
|
| 345 |
# Output projection
|
| 346 |
y = self.resid_dropout(self.c_proj(y))
|