pytorch - 💡(How to fix) Fix LazyBatchNorm2d complex64 input has inconsistent CPU and CUDA errors

Describe the bug

torch.nn.LazyBatchNorm2d fails inconsistently for complex64 input on CPU and CUDA.

CPU rejects the input with a direct unsupported-complex-dtype error, while CUDA reports a parameter/input dtype mismatch from cudnn_batch_norm.

The operation fails on both backends, but the CUDA error message is less direct and inconsistent with the CPU path.

Repro

import torch

x = torch.zeros((1, 3, 4, 4), dtype=torch.complex64)

try:
    print(torch.nn.LazyBatchNorm2d()(x).shape)
except Exception as e:
    print("CPU:", type(e).__name__, str(e))

try:
    print(torch.nn.LazyBatchNorm2d().cuda()(x.cuda()).shape)
except Exception as e:
    print("CUDA:", type(e).__name__, str(e))

Actual behavior

CPU: NotImplementedError "batch_norm" not implemented for 'ComplexFloat'
CUDA: RuntimeError Expected tensor for argument #1 'input' to have the same type as tensor for argument #2 'weight'; but type CUDAComplexFloatType does not equal torch.cuda.FloatTensor (while checking arguments for cudnn_batch_norm)

Expected behavior

CPU and CUDA should fail consistently, ideally with the same high-level unsupported-complex-dtype error.

For example, CUDA should also report that BatchNorm / LazyBatchNorm2d does not support ComplexFloat, rather than surfacing a lower-level parameter/input dtype mismatch.

Why this matters

This is a small backend consistency / error-reporting issue.

The current CUDA error looks like a parameter/input dtype mismatch, while the CPU error clearly says complex input is unsupported. A consistent error would make backend behavior easier to reason about and easier to test.

Environment

Python: 3.12.13
PyTorch: 2.10.0+cu128
CUDA: 12.8
GPU: Tesla T4

cc @csarofeen @ptrblck @eqy @nWEIdia @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @malfet @ezyang @anjali411 @dylanbespalko @nikitaved @amjames