fix: include torch compile flag

torch-ext/megablocks/layers.py

@@ -1073,6 +1073,7 @@ def get_device_mesh(model):
 
 
 class MegaBlocksMoeMLP(torch.nn.Module):
+    can_torch_compile: bool = True
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         moe_top_k = getattr(self.router, "top_k", 4)