Update modelopt/torch/quantization/plugins/huggingface.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com>

Author: cjluo-nv

modelopt/torch/quantization/plugins/huggingface.py

@@ -458,8 +458,10 @@ def _setup(self):
         elif hasattr(self, "experts") and hasattr(self.experts, "num_experts"):
             num_experts = self.experts.num_experts
 
-        self.expert_token_count = torch.zeros(
-            num_experts, dtype=torch.long, device=next(self.parameters()).device
+        self.register_buffer(
+            "expert_token_count",
+            torch.zeros(num_experts, dtype=torch.long, device=next(self.parameters()).device),
+            persistent=False,
         )
         self._count_expert_tokens = False
         self._moe_calib_experts_ratio = None