From 9a26d3fa23e1c933f87f15339e0d20f9b3fa6a97 Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Wed, 9 Sep 2026 07:46:27 -0700 Subject: [PATCH] Suppress new Pyrefly findings before upgrading to v1.2.0 PiperOrigin-RevId: 978542888 --- zapbench/models/nunet.py | 2 +- zapbench/video_forecasting/train.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/zapbench/models/nunet.py b/zapbench/models/nunet.py index 600a872..09a5e2a 100644 --- a/zapbench/models/nunet.py +++ b/zapbench/models/nunet.py @@ -476,7 +476,7 @@ def __call__( x = x_shard_fn(x) if d < len(resample_factors): - x = upsample(x, resample_factor) + x = upsample(x, resample_factor) # pyrefly: ignore[bad-argument-type] logging.info('Upsampled to shape %r', x.shape) # Upsample to superresolution diff --git a/zapbench/video_forecasting/train.py b/zapbench/video_forecasting/train.py index ec226c6..7dd2c6c 100644 --- a/zapbench/video_forecasting/train.py +++ b/zapbench/video_forecasting/train.py @@ -200,7 +200,7 @@ def loss_and_predictions( num_bins=num_bins, sigma_ratio=0.75, min_value=-0.25, max_value=1.5 ) target_shape = targets.shape - target_probs = jax.vmap(hlg_transform.transform_to_probs)(targets.flatten()) + target_probs = jax.vmap(hlg_transform.transform_to_probs)(targets.flatten()) # pyrefly: ignore[bad-argument-type] target_probs = target_probs.reshape(target_shape[:-1] + (num_bins,)) losses = optax.losses.softmax_cross_entropy(predictions, target_probs) if 'trace' in criterion: