From 3a072baffaea278cddc8b14798a45070ef174480 Mon Sep 17 00:00:00 2001
From: Giovanni Trezza <giotrezza298@gmail.com>
Date: Thu, 5 Mar 2026 11:44:21 +0100
Subject: [PATCH] Fix type checking

---
 src/reax/metrics/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/reax/metrics/utils.py b/src/reax/metrics/utils.py
index e91be0f..7249924 100644
--- a/src/reax/metrics/utils.py
+++ b/src/reax/metrics/utils.py
@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import TYPE_CHECKING, ClassVar, Optional, Protocol, TypeVar
+from typing import TYPE_CHECKING, ClassVar, Protocol, TypeVar
 
 import beartype
 import clu.internal.utils
@@ -20,7 +20,7 @@
 __all__ = tuple()
 
 M = TypeVar("M", bound=Metric)
-OptionalMask = Optional["reax.types.ArrayMask"]
+OptionalMask = types.ArrayMask | None
 
 
 class ReduceFn(Protocol):