From 1277188e39bb54e40a9e78f6a79f6dc59052ff22 Mon Sep 17 00:00:00 2001 From: Michal-Novomestsky Date: Sat, 22 Aug 2026 17:04:29 +1000 Subject: [PATCH] Only compile IsotropicModel.blocks --- walrus/models/isotropic_model.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/walrus/models/isotropic_model.py b/walrus/models/isotropic_model.py index 8d1974c..797038b 100755 --- a/walrus/models/isotropic_model.py +++ b/walrus/models/isotropic_model.py @@ -3,6 +3,7 @@ from operator import mul from typing import Callable, List, Optional +import warnings import numpy as np import torch import torch.nn as nn @@ -430,3 +431,8 @@ def forward( x = x.squeeze(-1) # Return T, B, C, H, [W], [D] return x # TODO - Return attention maps for debugging + + def compile(self, **kwargs) -> None: + warnings.warn("IsotropicModel.compile() currently only supports torch compilation for IsotropicModel.blocks. Note that torch.compile(IsotropicModel) will attempt to compile the entire forward call instead", UserWarning) + for blk in self.blocks: + blk.compile(**kwargs) \ No newline at end of file