[MUSA] Fix SwiGLU to use fused op (#6364)
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
1 parent
38cf9adc38
commit
b58feb1420
1 file changed
+1
-1
@@ -22,7 +22,7 @@ class SiluAndMul(CustomOp):
|
||||
return F.silu(x[..., :d]) * x[..., d:]
|
||||
|
||||
def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.device.type != "cuda":
|
||||
if x.device.type == "cpu":
|
||||
return self.forward_native(x)
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = x.shape[:-1] + (d,)
|
||||
|
||||
Reference in new issue
Block a user