From b5dc58373f84450dddcb914e7321ca4aa2f9be0c Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Fri, 19 Sep 2025 13:52:42 -0400 Subject: [PATCH] fix compile --- src/axolotl/kernels/moe/torch_grouped.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/axolotl/kernels/moe/torch_grouped.py b/src/axolotl/kernels/moe/torch_grouped.py index 68ff631bc..dd434d7a8 100644 --- a/src/axolotl/kernels/moe/torch_grouped.py +++ b/src/axolotl/kernels/moe/torch_grouped.py @@ -342,7 +342,7 @@ def moe_ffn_forward_grouped( counts_i32 = assignments.to(device=device, dtype=torch.int32) offsets = torch.cumsum(counts_i32, dim=0).to(dtype=torch.int32) - if offsets[-1].item() == 0: + if not torch.is_nonzero(offsets[-1]): zero = torch.zeros_like(x_flat) return zero.view(bsz, seqlen, hdim), router_logits mm_dtype = torch.bfloat16 if expert_dtype == torch.bfloat16 else expert_dtype