Fix conv2d bf16 crash on T4: iris/pde_ssm.py

iris/pde_ssm.py

@@ -183,9 +183,11 @@ class PDESSMBlock(nn.Module):
         x_2d = x.view(B, H, W, D).permute(0, 3, 1, 2)  # (B, D, H, W)
 
         # Spectral mixing (global) + Local conv (residual) + Token diff (anti-smoothing)
-        spectral_out = self.spectral(x_2d)
-        local_out = self.local_conv(x_2d)
-        mixed = spectral_out + local_out
+        # Run conv in float32 — grouped/1x1 convs lack bf16 cuDNN kernels on T4
+        with torch.amp.autocast(device_type='cuda', enabled=False):
+            spectral_out = self.spectral(x_2d.float())
+            local_out = self.local_conv(x_2d.float())
+        mixed = spectral_out.to(x.dtype) + local_out.to(x.dtype)
         mixed = self.token_diff(mixed)
 
         # Back to sequence format