Update dia2/runtime/context.py

dia2/runtime/context.py

@@ -46,7 +46,6 @@ def build_runtime(
     device_obj = torch.device(device)
     if device_obj.type == "cuda":
         cuda_matmul = torch.backends.cuda.matmul
-        cudnn_conv = torch.backends.cudnn.conv
         if hasattr(cuda_matmul, "fp32_precision"):
             cuda_matmul.fp32_precision = "tf32"
             with warnings.catch_warnings():
@@ -57,15 +56,22 @@ def build_runtime(
                 torch.backends.cuda.matmul.allow_tf32 = True
         else:  # pragma: no cover - compatibility with older PyTorch
             torch.backends.cuda.matmul.allow_tf32 = True
-        if hasattr(cudnn_conv, "fp32_precision"):
-            cudnn_conv.fp32_precision = "tf32"
-            with warnings.catch_warnings():
-                warnings.filterwarnings(
-                    "ignore",
-                    message="Please use the new API settings",
-                )
+
+        # Handle cuDNN conv TF32 settings (check if conv attribute exists first)
+        if hasattr(torch.backends.cudnn, "conv"):
+            cudnn_conv = torch.backends.cudnn.conv
+            if hasattr(cudnn_conv, "fp32_precision"):
+                cudnn_conv.fp32_precision = "tf32"
+                with warnings.catch_warnings():
+                    warnings.filterwarnings(
+                        "ignore",
+                        message="Please use the new API settings",
+                    )
+                    torch.backends.cudnn.allow_tf32 = True
+            else:
                 torch.backends.cudnn.allow_tf32 = True
-        else:  # pragma: no cover
+        else:
+            # For older PyTorch versions without the conv attribute
             torch.backends.cudnn.allow_tf32 = True
     precision = resolve_precision(dtype_pref, device_obj)
     config = load_config(config_path)