diff --git a/codes/models/archs/SwitchedResidualGenerator_arch.py b/codes/models/archs/SwitchedResidualGenerator_arch.py
index c214819d..ea0a167c 100644
--- a/codes/models/archs/SwitchedResidualGenerator_arch.py
+++ b/codes/models/archs/SwitchedResidualGenerator_arch.py
@@ -123,7 +123,7 @@ class ConfigurableSwitchComputer(nn.Module):
         self.noise_scale = nn.Parameter(torch.full((1,), float(1e-3)))
 
         # And the switch itself, including learned scalars
-        self.switch = BareConvSwitch(initial_temperature=init_temp)
+        self.switch = BareConvSwitch(initial_temperature=init_temp, attention_norm=AttentionNorm(transform_count, accumulator_size=128))
         self.switch_scale = nn.Parameter(torch.full((1,), float(1)))
         self.post_switch_conv = ConvBnLelu(base_filters, base_filters, norm=False, bias=True)
         # The post_switch_conv gets a low scale initially. The network can decide to magnify it (or not)