From c44895c36c0bb95a7371ff52b1a8ec8b99cf58e9 Mon Sep 17 00:00:00 2001 From: Alex-Brooks Date: Thu, 12 Feb 2026 17:26:47 +0000 Subject: [PATCH] fix missing import Signed-off-by: Alex-Brooks --- aiu_fms_testing_utils/utils/paged.py | 1 + 1 file changed, 1 insertion(+) diff --git a/aiu_fms_testing_utils/utils/paged.py b/aiu_fms_testing_utils/utils/paged.py index 49807dc1..377280af 100644 --- a/aiu_fms_testing_utils/utils/paged.py +++ b/aiu_fms_testing_utils/utils/paged.py @@ -4,6 +4,7 @@ import time from typing import Any, Callable, List, MutableMapping, Optional, Tuple, Union import torch +import torch.nn.functional as F import fms.utils.spyre.paged # noqa from aiu_fms_testing_utils.utils import get_pad_size