From 4ba0663d68a618b6608da630a1cc18eb117d6fe2 Mon Sep 17 00:00:00 2001 From: V-MoE Authors Date: Wed, 5 Aug 2026 12:39:21 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 959811586 --- vmoe/checkpoints/base.py | 2 +- vmoe/train/optimizer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/vmoe/checkpoints/base.py b/vmoe/checkpoints/base.py index 0ceb814..7877c90 100644 --- a/vmoe/checkpoints/base.py +++ b/vmoe/checkpoints/base.py @@ -118,7 +118,7 @@ def _parse_step_and_suffix_or_error(filepath): sorted(steps_and_suffixes - incomplete_steps_and_suffixes, reverse=decreasing), lambda x: x[0]): - if not suffixes.difference(set(x[1] for x in group)): # pyrefly: ignore[missing-attribute] + if not suffixes.difference(set(x[1] for x in group)): yield step diff --git a/vmoe/train/optimizer.py b/vmoe/train/optimizer.py index abb688e..431d461 100644 --- a/vmoe/train/optimizer.py +++ b/vmoe/train/optimizer.py @@ -166,7 +166,7 @@ def frozen_fn(params: optax.Params) -> PyTree: flatparams = flax.traverse_util.flatten_dict( flax.serialization.to_state_dict(params), sep='/') output = { - key: search_true if pattern.search(key) else search_false # pyrefly: ignore[missing-attribute] + key: search_true if pattern.search(key) else search_false for key, value in flatparams.items() } return flax.serialization.from_state_dict(