[model] fix: import Gemma4 checkpoint from its defining module

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/models/gemma/modeling_gemma4.py

@@ -39,6 +39,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from megatron.core import parallel_state, tensor_parallel
+from megatron.core.extensions.transformer_engine import te_checkpoint
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.models.backends import LocalSpecProvider
@@ -49,7 +50,7 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
-from megatron.core.transformer.moe.moe_layer import MoELayer, te_checkpoint
+from megatron.core.transformer.moe.moe_layer import MoELayer
 from megatron.core.transformer.moe.router import TopKRouter
 from megatron.core.transformer.spec_utils import ModuleSpec, get_submodules
 from megatron.core.transformer.transformer_config import TransformerConfig