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

iris/model.py

@@ -17,7 +17,11 @@ class Patchify(nn.Module):
     def forward(self, z):
         B, C, H, W = z.shape
         p = self.patch_size
-        z = self.dw_conv(z)
+        orig_dtype = z.dtype
+        # Run grouped conv in float32 — cuDNN lacks bf16 kernels for grouped convs on T4
+        with torch.amp.autocast(device_type='cuda', enabled=False):
+            z = self.dw_conv(z.float())
+        z = z.to(orig_dtype)
         H_tok, W_tok = H // p, W // p
         z = z.view(B, C, H_tok, p, W_tok, p).permute(0, 2, 4, 1, 3, 5).reshape(B, H_tok * W_tok, C * p * p)
         return self.proj(z), H_tok, W_tok
@@ -37,7 +41,11 @@ class Unpatchify(nn.Module):
         C = self.out_channels
         z = self.proj(tokens).view(B, H_tok, W_tok, C, p, p)
         z = z.permute(0, 3, 1, 4, 2, 5).reshape(B, C, H_tok * p, W_tok * p)
-        return self.dw_conv(z)
+        # Run grouped conv in float32 — cuDNN lacks bf16 kernels for grouped convs on T4
+        orig_dtype = z.dtype
+        with torch.amp.autocast(device_type='cuda', enabled=False):
+            z = self.dw_conv(z.float())
+        return z.to(orig_dtype)
 
 
 class TinyDecoder(nn.Module):
@@ -55,10 +63,14 @@ class TinyDecoder(nn.Module):
         self.final = nn.Conv2d(out_channels, out_channels, 1, bias=True)
 
     def forward(self, z):
-        x = z
-        for stage in self.stages:
-            x = stage(x)
-        return torch.tanh(self.final(x))
+        # Run decoder convs in float32 — cuDNN lacks bf16 kernels on T4
+        orig_dtype = z.dtype
+        with torch.amp.autocast(device_type='cuda', enabled=False):
+            x = z.float()
+            for stage in self.stages:
+                x = stage(x)
+            x = torch.tanh(self.final(x))
+        return x.to(orig_dtype)
 
 
 class IRIS(nn.Module):