You write custom CUDA kernels to replace the pytorch operators in the given architecture to get speedups. 

    You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining matmul+relu), or algorithmic changes (such as online softmax). You are only limited by your imagination.


        Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is: 

        ``` 

        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, a, b):
        return a + b


def get_inputs():
    # randomly generate input tensors based on the model architecture
    a = torch.randn(1, 128).cuda()
    b = torch.randn(1, 128).cuda()
    return [a, b]


def get_init_inputs():
    # randomly generate tensors required for initialization based on the model architecture
    return []

        ``` 

        The example new arch with custom CUDA kernels looks like this:
        ```
        # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Define the custom CUDA kernel for element-wise addition
elementwise_add_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void elementwise_add_kernel(const float* a, const float* b, float* out, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = a[idx] + b[idx];
    }
}

torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b) {
    auto size = a.numel();
    auto out = torch::zeros_like(a);

    const int block_size = 256;
    const int num_blocks = (size + block_size - 1) / block_size;

    elementwise_add_kernel<<<num_blocks, block_size>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), size);

    return out;
}
"""

elementwise_add_cpp_source = (
    "torch::Tensor elementwise_add_cuda(torch::Tensor a, torch::Tensor b);"
)

# Compile the inline CUDA code for element-wise addition
elementwise_add = load_inline(
    name="elementwise_add",
    cpp_sources=elementwise_add_cpp_source,
    cuda_sources=elementwise_add_source,
    functions=["elementwise_add_cuda"],
    verbose=True,
    extra_cflags=[""],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.elementwise_add = elementwise_add

    def forward(self, a, b):
        return self.elementwise_add.elementwise_add_cuda(a, b)

        ``` 

        
    You are given the following architecture: 

    ```
    # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, scales the output, applies batch normalization, 
    and then performs global average pooling. 
    """
    def __init__(self, in_channels, out_channels, kernel_size, scale_factor, eps=1e-5, momentum=0.1):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = scale_factor
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))

    def forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

batch_size = 16
in_channels = 64
out_channels = 128
depth, height, width = 16, 32, 32
kernel_size = 5
scale_factor = 2.0

def get_inputs():
    return [torch.rand(batch_size, in_channels, depth, height, width)]

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, scale_factor]
    ```
    
Optimize the architecture named Model with custom CUDA operators! Name your optimized output architecture ModelNew. Output the new code in codeblocks. Please generate real code, NOT pseudocode, make sure the code compiles and is fully functional. Just output the new model code, no other text, and NO testing code! 



This is turn 3 of 6 for this task. Your previous solution (turn 2) was:
```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Algorithmic rewrite of: ConvTranspose3d -> scale -> BatchNorm3d -> global avg pool.
#
# The network ends in a *global* average pool and BatchNorm is a per-channel
# affine transform, so the result only depends on the per-(batch, channel)
# spatial sums and sums-of-squares of the transposed-convolution output.
#
# For stride-1 / pad-0 ConvTranspose3d, the output is the *full* linear
# convolution of x and w, so by the convolution theorem and Parseval:
#     sum_z y[c,z]  = Y[c,0]                     (DC Fourier component)
#     sum_z y[c,z]^2 = (1/N) * sum_f |Y[c,f]|^2  (Parseval)
# with Y[c,f] = sum_i X[i,f] * W[i,c,f] (a small complex GEMM per frequency).
# This avoids materializing the ~200MB / ~850 GFLOP convolution entirely:
#   cost ~ one rfftn of the input + one memory-bound complex GEMM.
# The pointwise |Y|^2 reduction is fused into the GEMM kernel, and a final
# tiny kernel applies conv bias, scale, BatchNorm (batch or running stats)
# and emits the globally pooled result.
# ---------------------------------------------------------------------------

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>

#define BLOCK_SIZE 256

// ------------------------------ helpers -----------------------------------
__device__ __forceinline__ float warp_reduce_sum(float v) {
    for (int off = 16; off > 0; off >>= 1)
        v += __shfl_down_sync(0xffffffffu, v, off);
    return v;
}

// Per-(batch, channel) spatial sums of a contiguous NCDHW tensor.
// One block per (b, c).
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    int S) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long base = ((long)b * (long)gridDim.x + (long)c) * (long)S;
    const float* xp = x + base;
    float acc = 0.f;
    if ((S & 3) == 0) {
        const float4* xp4 = reinterpret_cast<const float4*>(xp);
        int S4 = S >> 2;
        for (int s = threadIdx.x; s < S4; s += BLOCK_SIZE) {
            float4 v = xp4[s];
            acc += (v.x + v.y) + (v.z + v.w);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) acc += xp[s];
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    acc = warp_reduce_sum(acc);
    int lane = threadIdx.x & 31, wid = threadIdx.x >> 5;
    if (lane == 0) sh[wid] = acc;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < BLOCK_SIZE / 32) ? sh[lane] : 0.f;
        v = warp_reduce_sum(v);
        if (lane == 0) sums[(long)b * (long)gridDim.x + (long)c] = v;
    }
}

torch::Tensor spatial_sums(torch::Tensor x) {
    auto xc = x.contiguous();
    int64_t B = xc.size(0), C = xc.size(1), S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    auto sums = torch::empty({B, C}, xc.options());
    if (B * C == 0) return sums;
    TORCH_CHECK(B <= 65535, "batch too large");
    dim3 grid((unsigned)C, (unsigned)B);
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sums_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(), (int)S);
    return sums;
}

// -------------------- fused frequency GEMM + Parseval ---------------------
// X: [F, B, I] complex (float2), W: [F, I, O] complex.
// For every frequency f the block computes ACC[b,o] = sum_i X[f,b,i]*W[f,i,o]
// (complex), then folds the Parseval weight:
//   block f==0            -> sums[b,o]  = Re(ACC[b,o])
//   block with 0<f<dc_blk -> sumsq[b,o] += 2*|ACC|^2 (interior frequencies)
//   block f==dc_blk       -> sumsq[b,o] += w_nyq*|ACC|^2 (DC/Nyquist planes)
// Grid: (dc_blk+1, num_tiles), block: 256 threads, 8 (b,o) pairs each.
// Requires B*O <= 2048, I <= 1024, O % 32 == 0.
__global__ void freq_gemm_parseval_kernel(const float2* __restrict__ X,
                                          const float2* __restrict__ W,
                                          float* __restrict__ sums,
                                          float* __restrict__ sumsq,
                                          int F, int B, int I, int O,
                                          int dc_blk, float nyq_w,
                                          float inv_n) {
    int f = blockIdx.x;
    int tile = blockIdx.y;
    int pbase = tile * 2048;

    __shared__ float2 sX[16 * 64];        // [B][I] tile
    __shared__ float2 sW[64 * 32];        // [I][32] chunk
    __shared__ float sDC[2048];           // DC accumulators (f == 0 only)

    const int BI = B * I;
    for (int t = threadIdx.x; t < BI; t += BLOCK_SIZE)
        sX[t] = X[(long)f * BI + t];
    if (f == 0) {
        for (int t = threadIdx.x; t < 2048; t += BLOCK_SIZE) sDC[t] = 0.f;
    }
    __syncthreads();

    float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f;
    float acc4 = 0.f, acc5 = 0.f, acc6 = 0.f, acc7 = 0.f;
    float bw[8];
    #pragma unroll
    for (int k = 0; k < 8; ++k) {
        int p = pbase + k * BLOCK_SIZE + threadIdx.x;
        int b = (p < B * O) ? (p / O) : 0;
        bw[k] = (float)b * (float)I;
    }

    for (int oc = 0; oc < O; oc += 32) {
        for (int t = threadIdx.x; t < (I << 5); t += BLOCK_SIZE)
            sW[t] = W[(long)f * I * O + (long)(t >> 5) * O + oc + (t & 31)];
        __syncthreads();
        #pragma unroll
        for (int k = 0; k < 8; ++k) {
            int p = pbase + k * BLOCK_SIZE + threadIdx.x;
            if (p >= B * O) break;
            int o = p - (int)bw[k] / (float)I * 0;  // placeholder (unused)
            o = p % O;
            float2 a = make_float2(0.f, 0.f);
            const float2* xrow = sX + (int)bw[k];
            const float2* wcol = sW + (o - oc);
            for (int i = 0; i < I; ++i) {
                float2 xv = __ldg(xrow + i);
                float2 wv = wcol[i << 5];
                a.x += xv.x * wv.x - xv.y * wv.y;
                a.y += xv.x * wv.y + xv.y * wv.x;
            }
            float pw = a.x * a.x + a.y * a.y;
            if (f == 0) {
                atomicAdd(&sDC[k * BLOCK_SIZE + threadIdx.x], a.x);
                if (oc == 0) acc0 += 0.f;  // keep structure simple
            }
            switch (k) {
                case 0: acc0 += pw; break; case 1: acc1 += pw; break;
                case 2: acc2 += pw; break; case 3: acc3 += pw; break;
                case 4: acc4 += pw; break; case 5: acc5 += pw; break;
                case 6: acc6 += pw; break; default: acc7 += pw; break;
            }
        }
        __syncthreads();
    }

    if (f == 0) {
        __syncthreads();
        for (int t = threadIdx.x; t < 2048; t += BLOCK_SIZE) {
            if (pbase + t < B * O) sums[pbase + t] = sDC[t];
        }
    }

    float wgt;
    if (f == 0 || f == dc_blk) wgt = nyq_w * inv_n;
    else                       wgt = 2.0f * inv_n;
    #pragma unroll
    for (int k = 0; k < 8; ++k) {
        int p = pbase + k * BLOCK_SIZE + threadIdx.x;
        if (p >= B * O) break;
        float v;
        switch (k) {
            case 0: v = acc0; break; case 1: v = acc1; break;
            case 2: v = acc2; break; case 3: v = acc3; break;
            case 4: v = acc4; break; case 5: v = acc5; break;
            case 6: v = acc6; break; default: v = acc7; break;
        }
        atomicAdd(&sumsq[p], v * wgt);
    }
}

std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t dc_blk, double nyq_w,
                                              double inv_n) {
    int F = X.size(0);
    auto opts = X.options().dtype(at::kFloat);
    auto sums  = torch::empty({B * O}, opts);
    auto sumsq = torch::zeros({B * O}, opts);
    int BO = B * (int)O;
    int tiles = (BO + 2047) / 2048;
    dim3 grid((unsigned)(dc_blk + 1), (unsigned)tiles);
    auto stream = at::cuda::getCurrentCUDAStream();
    freq_gemm_parseval_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        reinterpret_cast<const float2*>(X.data_ptr<float>()),
        reinterpret_cast<const float2*>(W.data_ptr<float>()),
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        F, (int)B, (int)I, (int)O, (int)dc_blk, (float)nyq_w, (float)inv_n);
    return {sums, sumsq};
}

// ------------------------------ finalize ----------------------------------
// Applies conv-bias correction to the moments, then scale + BatchNorm and
// emits the globally pooled result. Training mode uses batch statistics and
// updates running stats exactly like PyTorch (biased var for normalize,
// unbiased for running). sum0: raw spatial sums (DC), sumsq0: raw sums of
// squares (Parseval). y = conv(x) + bias; value = scale * y.
__global__ void finalize_kernel(const float* __restrict__ sum0,
                                const float* __restrict__ sumsq0,
                                const float* __restrict__ bias,
                                float* __restrict__ rmean,
                                float* __restrict__ rvar,
                                const float* __restrict__ gamma,
                                const float* __restrict__ beta,
                                float* __restrict__ out,
                                int B, int C, float S,
                                float scale, float momentum,
                                float eps, int training) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * C) return;
    int b = idx / C;
    int c = idx - b * C;
    float bc = (bias != nullptr) ? bias[c] : 0.f;
    float s_raw  = sum0[idx];
    float s2_raw = sumsq0[idx];
    float su  = s_raw + bc * S;
    float s2u = s2_raw + 2.f * bc * s_raw + bc * bc * S;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            float sr = sum0[(long)bb * C + c];
            float s2r = sumsq0[(long)bb * C + c];
            s  += sr + bc * S;
            s2 += s2r + 2.f * bc * sr + bc * bc * S;
        }
        s  *= scale;
        s2 *= scale * scale;
        mean = s / N;
        float var = s2 / N - mean * mean;
        var = fmaxf(var, 0.f);
        invstd = rsqrtf(var + eps);
        if (b == 0 && rmean != nullptr && momentum >= 0.f) {
            float unbiased = var * (N / (N - 1.f));
            rmean[c] = (1.f - momentum) * rmean[c] + momentum * mean;
            rvar[c]  = (1.f - momentum) * rvar[c]  + momentum * unbiased;
        }
    } else {
        mean = rmean[c];
        invstd = rsqrtf(rvar[c] + eps);
    }
    float g  = (gamma != nullptr) ? gamma[c] : 1.f;
    float be = (beta  != nullptr) ? beta[c]  : 0.f;
    float mbc = su * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, torch::Tensor rmean,
                       torch::Tensor rvar, torch::Tensor gamma,
                       torch::Tensor beta, double scale, double S,
                       double momentum, double eps, bool training) {
    int B = sums.size(0), C = sums.size(1);
    auto out = torch::empty({B, C}, sums.options());
    int total = B * C;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.numel() ? rmean.data_ptr<float>() : nullptr,
        rvar.numel()  ? rvar.data_ptr<float>()  : nullptr,
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, C, (float)S, (float)scale,
        (float)momentum, (float)eps, training ? 1 : 0);
    return out;
}

// Eval closed form: per-channel sums of the conv output via
// sum(out[b,o]) = sum_i Sx[b,i] * Sw[i,o] + bias[o]*S_out, then scale+BN+pool.
__global__ void eval_finalize_kernel(const float* __restrict__ Sx,
                                     const float* __restrict__ Sw,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ rmean,
                                     const float* __restrict__ rvar,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ out,
                                     int B, int In, int Out,
                                     float scale, float inv_Sout, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * Out) return;
    int b = idx / Out;
    int o = idx - b * Out;
    const float* sx = Sx + (long)b * In;
    float acc = 0.f;
    for (int i = 0; i < In; ++i) acc += sx[i] * Sw[(long)i * Out + o];
    float mean = acc * inv_Sout;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float g  = (gamma != nullptr) ? gamma[o] : 1.f;
    float be = (beta  != nullptr) ? beta[o]  : 0.f;
    out[idx] = (mean - rmean[o]) * rsqrtf(rvar[o] + eps) * g + be;
}

torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta, double scale,
                            double inv_Sout, double eps) {
    int B = Sx.size(0), In = Sx.size(1), Out = Sw.size(1);
    auto out = torch::empty({B, Out}, Sx.options());
    int total = B * Out;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    eval_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, In, Out,
        (float)scale, (float)inv_Sout, (float)eps);
    return out;
}
"""

_cpp_src = r"""
#include <vector>
torch::Tensor spatial_sums(torch::Tensor x);
std::vector<torch::Tensor> freq_gemm_parseval(torch::Tensor X, torch::Tensor W,
                                              int64_t B, int64_t I, int64_t O,
                                              int64_t dc_blk, double nyq_w,
                                              double inv_n);
torch::Tensor finalize(torch::Tensor sums, torch::Tensor sumsq,
                       torch::Tensor bias, torch::Tensor rmean,
                       torch::Tensor rvar, torch::Tensor gamma,
                       torch::Tensor beta, double scale, double S,
                       double momentum, double eps, bool training);
torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta, double scale,
                            double inv_Sout, double eps);
"""

_ext = load_inline(
    name="fused_fft_convtranspose3d_bn_globalpool",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["spatial_sums", "freq_gemm_parseval", "finalize", "eval_finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Drop-in optimized replacement for Model (identical submodules/state dict).

    Train mode: the ConvTranspose3d is evaluated in the Fourier domain
    (full convolution == pointwise product of spectra).  Only the DC
    component (spatial sums) and the Parseval power (sums of squares) of the
    output spectrum are computed -- fused into one custom CUDA kernel -- and a
    tiny finalize kernel applies bias, scale, BatchNorm (with running-stat
    update) and global average pooling.  The huge conv output tensor and the
    ~850 GFLOP direct convolution are never materialized.

    Eval mode: spatial sums of the conv output are computed in closed form
    from per-channel input sums and kernel sums (also exact), so no
    convolution runs at all.
    """

    def __init__(self, in_channels, out_channels, kernel_size, scale_factor,
                 eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = float(scale_factor)
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        ct = self.conv_transpose
        self._simple_conv = (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )
        self._w_spec = None   # cached weight spectrum [F, In, Out] (complex64)
        self._w_key = None

    # ------------------------------------------------------------------ util
    def _reference_forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

    def _out_spatial(self, x):
        ct = self.conv_transpose
        D, H, W = x.shape[2], x.shape[3], x.shape[4]
        k = ct.kernel_size
        return (D + k[0] - 1, H + k[1] - 1, W + k[2] - 1)

    def _get_weight_spectrum(self, s):
        ct = self.conv_transpose
        w = ct.weight
        key = (w._version, w.data_ptr(), s)
        if self._w_key != key or self._w_spec is None:
            wp = w.detach().permute(2, 3, 4, 0, 1).contiguous()  # [kd,kh,kw,In,Out]
            spec = torch.fft.rfftn(wp, s=s, dim=(0, 1, 2))       # [D,H,Wh,In,Out]
            F = spec.shape[0] * spec.shape[1] * spec.shape[2]
            self._w_spec = spec.reshape(F, spec.shape[3], spec.shape[4]).contiguous()
            self._w_key = key
        return self._w_spec

    # ---------------------------------------------------------------- forward
    def forward(self, x):
        bn = self.batch_norm
        ct = self.conv_transpose

        if not (x.is_cuda and x.dtype == torch.float32):
            return self._reference_forward(x)
        if bn.training and bn.momentum is None:
            return self._reference_forward(x)  # cumulative moving average
        if (not bn.training) and bn.track_running_stats and (
            bn.running_mean is None or not bn.running_mean.is_cuda
        ):
            return self._reference_forward(x)

        B = x.shape[0]
        In, Out = ct.in_channels, ct.out_channels
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)
        bias = ct.bias if ct.bias is not None else x.new_empty(0)

        if bn.track_running_stats:
            rmean, rvar = bn.running_mean, bn.running_var
        else:
            rmean = rvar = x.new_empty(0)
        momentum = bn.momentum if bn.momentum is not None else -1.0

        if not self._simple_conv:
            # Generic fallback: run the conv, fuse the rest.
            y = ct(x)
            C = y.shape[1]
            S = y.shape[2] * y.shape[3] * y.shape[4]
            yv = (y * self.scale_factor).contiguous()
            sums = _ext.spatial_sums(yv)
            sumsq = _ext.spatial_sums(yv * yv)
            if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1
            out = _ext.finalize(sums, sumsq, x.new_empty(0), rmean, rvar,
                                gamma, beta, 1.0, float(S), momentum,
                                bn.eps, bn.training)
            return out.view(B, C, 1, 1, 1)

        Do, Ho, Wo = self._out_spatial(x)
        S = Do * Ho * Wo
        Wh = Wo // 2 + 1
        F = Do * Ho * Wh

        if (not bn.training) and bn.track_running_stats:
            # ---- closed-form eval path: no convolution ---------------------
            Sx = _ext.spatial_sums(x)               # [B, In]
            Sw = _ext.spatial_sums(ct.weight)       # [In, Out]
            out = _ext.eval_finalize(Sx, Sw, bias, rmean, rvar, gamma, beta,
                                     self.scale_factor, 1.0 / float(S), bn.eps)
            return out.view(B, Out, 1, 1, 1)

        # -------------------- FFT path (train / no-track eval) --------------
        # Custom fused kernel constraints.
        if not (B * Out <= 2048 and In <= 64 and Out % 32 == 0):
            return self._reference_forward(x)

        s = (Do, Ho, Wo)
        xp = x.detach().permute(2, 3, 4, 0, 1)                     # [D,H,W,B,In]
        Xf = torch.fft.rfftn(xp, s=s, dim=(0, 1, 2))               # [D,H,Wh,B,In]
        Xf = Xf.reshape(F, B, In)
        Wf = self._get_weight_spectrum(s)                          # [F, In, Out]

        dc_blk = Ho * Wh                                           # d == 0 plane
        nyq_w = 2.0 if (Wo % 2 == 1) else 1.0
        sums, sumsq = _ext.freq_gemm_parseval(
            Xf, Wf, B, In, Out, dc_blk, nyq_w, 1.0 / float(S))

        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.finalize(sums.view(B, Out), sumsq.view(B, Out), bias,
                            rmean, rvar, gamma, beta, self.scale_factor,
                            float(S), momentum, bn.eps, bn.training)
        return out.view(B, Out, 1, 1, 1)
```
Result of evaluating it:
The solution compiled but failed at run time. Error output:
```
expected scalar type Float but found ComplexFloat
```

The fastest correct solution so far (1.085x speedup) is:
```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# ---------------------------------------------------------------------------
# Custom CUDA kernels.
#
# Key insight: the network ends in a *global* average pool, and BatchNorm is a
# per-channel affine transform. Therefore the final result only depends on the
# per-channel spatial sums (and sums of squares, for train-mode batch stats)
# of the transposed-convolution output -- the full output tensor never needs
# to be materialized or re-read.
#
# In eval mode (fixed running stats) we can go one step further: the spatial
# sum of a stride-1/pad-0 ConvTranspose3d output decomposes exactly as
#     sum(out[b,o]) = sum_i sum(x[b,i]) * sum(w[i,o]) + bias[o] * S_out
# so the entire 500+ GFLOP convolution collapses into two cheap reductions
# plus a tiny 16x64x128 contraction, fused with scale + BN + pooling.
#
# In training mode we run the (unavoidable) convolution and fuse
# scale + batch-stat computation + running-stat update + BN + global pooling
# into a pair of lightweight reduction kernels.
# ---------------------------------------------------------------------------

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>

#define BLOCK_SIZE 256

// Per-(batch, channel) spatial sums (and optionally sums of squares) of a
// contiguous NCDHW tensor. One block per (b, c) pair.
__global__ void spatial_sums_kernel(const float* __restrict__ x,
                                    float* __restrict__ sums,
                                    float* __restrict__ sumsq,
                                    int S, int compute_sq) {
    int c = blockIdx.x;
    int b = blockIdx.y;
    long base = ((long)b * (long)gridDim.x + (long)c) * (long)S;
    const float* xp = x + base;

    float acc = 0.f, acc2 = 0.f;
    if ((S & 3) == 0) {
        const float4* xp4 = reinterpret_cast<const float4*>(xp);
        int S4 = S >> 2;
        for (int s = threadIdx.x; s < S4; s += BLOCK_SIZE) {
            float4 v = xp4[s];
            acc += (v.x + v.y) + (v.z + v.w);
            if (compute_sq) acc2 += (v.x * v.x + v.y * v.y) + (v.z * v.z + v.w * v.w);
        }
    } else {
        for (int s = threadIdx.x; s < S; s += BLOCK_SIZE) {
            float v = xp[s];
            acc += v;
            if (compute_sq) acc2 += v * v;
        }
    }

    // warp reduction
    for (int off = 16; off > 0; off >>= 1) {
        acc += __shfl_down_sync(0xffffffffu, acc, off);
        if (compute_sq) acc2 += __shfl_down_sync(0xffffffffu, acc2, off);
    }
    __shared__ float sh[BLOCK_SIZE / 32];
    __shared__ float sh2[BLOCK_SIZE / 32];
    int lane = threadIdx.x & 31;
    int wid = threadIdx.x >> 5;
    if (lane == 0) {
        sh[wid] = acc;
        if (compute_sq) sh2[wid] = acc2;
    }
    __syncthreads();
    if (wid == 0) {
        int nw = BLOCK_SIZE / 32;
        float v  = (lane < nw) ? sh[lane] : 0.f;
        float v2 = (lane < nw && compute_sq) ? sh2[lane] : 0.f;
        for (int off = 16; off > 0; off >>= 1) {
            v += __shfl_down_sync(0xffffffffu, v, off);
            if (compute_sq) v2 += __shfl_down_sync(0xffffffffu, v2, off);
        }
        if (lane == 0) {
            long idx = (long)b * (long)gridDim.x + (long)c;
            sums[idx] = v;
            if (compute_sq) sumsq[idx] = v2;
        }
    }
}

std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == at::kFloat, "x must be float32");
    auto xc = x.contiguous();
    int64_t B = xc.size(0);
    int64_t C = xc.size(1);
    int64_t S = 1;
    for (int i = 2; i < xc.dim(); ++i) S *= xc.size(i);
    TORCH_CHECK(B <= 65535, "batch too large");
    auto sums  = torch::empty({B, C}, xc.options());
    auto sumsq = compute_sq ? torch::empty({B, C}, xc.options())
                            : torch::empty({0}, xc.options());
    if (B * C == 0) return {sums, sumsq};
    dim3 grid((unsigned)C, (unsigned)B);
    auto stream = at::cuda::getCurrentCUDAStream();
    spatial_sums_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        xc.data_ptr<float>(), sums.data_ptr<float>(),
        compute_sq ? sumsq.data_ptr<float>() : nullptr,
        (int)S, compute_sq ? 1 : 0);
    return {sums, sumsq};
}

// Eval-mode fused kernel: contracts input sums with kernel sums, applies
// conv bias, scale, BatchNorm (running stats) and emits the globally pooled
// result. Sw is [In, Out] (ConvTranspose weight layout is (In, Out, kD,kH,kW)).
__global__ void eval_finalize_kernel(const float* __restrict__ Sx,
                                     const float* __restrict__ Sw,
                                     const float* __restrict__ bias,
                                     const float* __restrict__ rmean,
                                     const float* __restrict__ rvar,
                                     const float* __restrict__ gamma,
                                     const float* __restrict__ beta,
                                     float* __restrict__ out,
                                     int B, int In, int Out,
                                     float scale, float inv_Sout, float eps) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * Out) return;
    int b = idx / Out;
    int o = idx - b * Out;
    const float* sx = Sx + (long)b * In;
    float acc = 0.f;
    for (int i = 0; i < In; ++i) acc += sx[i] * Sw[(long)i * Out + o];
    float mean = acc * inv_Sout;
    if (bias != nullptr) mean += bias[o];
    mean *= scale;
    float g  = (gamma != nullptr) ? gamma[o] : 1.f;
    float be = (beta  != nullptr) ? beta[o]  : 0.f;
    out[idx] = (mean - rmean[o]) * rsqrtf(rvar[o] + eps) * g + be;
}

torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps) {
    int B = Sx.size(0), In = Sx.size(1), Out = Sw.size(1);
    auto out = torch::empty({B, Out}, Sx.options());
    int total = B * Out;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    eval_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        Sx.data_ptr<float>(), Sw.data_ptr<float>(),
        bias.numel()  ? bias.data_ptr<float>()  : nullptr,
        rmean.data_ptr<float>(), rvar.data_ptr<float>(),
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, In, Out,
        (float)scale, (float)inv_Sout, (float)eps);
    return out;
}

// Train-mode (or generic) fused kernel: given per-(b,c) sums and sums of
// squares of the convolution output, compute batch statistics, update the
// running statistics in place (exactly like PyTorch: biased variance for
// normalization, unbiased for the running average), then apply the affine
// transform and emit the globally pooled result. With training=0 the running
// statistics are used directly (eval fallback for non-unit conv params).
__global__ void train_finalize_kernel(const float* __restrict__ sums,
                                      const float* __restrict__ sumsq,
                                      float* __restrict__ rmean,
                                      float* __restrict__ rvar,
                                      const float* __restrict__ gamma,
                                      const float* __restrict__ beta,
                                      float* __restrict__ out,
                                      int B, int C, float S,
                                      float scale, float momentum,
                                      float eps, int training) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= B * C) return;
    int b = idx / C;
    int c = idx - b * C;
    float mean, invstd;
    float N = S * (float)B;
    if (training) {
        float s = 0.f, s2 = 0.f;
        for (int bb = 0; bb < B; ++bb) {
            s  += sums[(long)bb * C + c];
            s2 += sumsq[(long)bb * C + c];
        }
        s  *= scale;
        s2 *= scale * scale;
        mean = s / N;
        float var = s2 / N - mean * mean;
        var = fmaxf(var, 0.f);
        invstd = rsqrtf(var + eps);
        if (b == 0 && rmean != nullptr && momentum >= 0.f) {
            float unbiased = var * (N / (N - 1.f));
            rmean[c] = (1.f - momentum) * rmean[c] + momentum * mean;
            rvar[c]  = (1.f - momentum) * rvar[c]  + momentum * unbiased;
        }
    } else {
        mean = rmean[c];
        invstd = rsqrtf(rvar[c] + eps);
    }
    float g  = (gamma != nullptr) ? gamma[c] : 1.f;
    float be = (beta  != nullptr) ? beta[c]  : 0.f;
    float mbc = sums[(long)b * C + c] * (scale / S);
    out[idx] = (mbc - mean) * invstd * g + be;
}

torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training) {
    int B = sums.size(0), C = sums.size(1);
    auto out = torch::empty({B, C}, sums.options());
    int total = B * C;
    if (total == 0) return out;
    auto stream = at::cuda::getCurrentCUDAStream();
    train_finalize_kernel<<<(total + 255) / 256, 256, 0, stream>>>(
        sums.data_ptr<float>(), sumsq.data_ptr<float>(),
        rmean.numel() ? rmean.data_ptr<float>() : nullptr,
        rvar.numel()  ? rvar.data_ptr<float>()  : nullptr,
        gamma.numel() ? gamma.data_ptr<float>() : nullptr,
        beta.numel()  ? beta.data_ptr<float>()  : nullptr,
        out.data_ptr<float>(), B, C, (float)S, (float)scale,
        (float)momentum, (float)eps, training ? 1 : 0);
    return out;
}
"""

_cpp_src = r"""
#include <vector>
std::vector<torch::Tensor> spatial_sums(torch::Tensor x, bool compute_sq);
torch::Tensor eval_finalize(torch::Tensor Sx, torch::Tensor Sw,
                            torch::Tensor bias, torch::Tensor rmean,
                            torch::Tensor rvar, torch::Tensor gamma,
                            torch::Tensor beta,
                            double scale, double inv_Sout, double eps);
torch::Tensor train_finalize(torch::Tensor sums, torch::Tensor sumsq,
                             torch::Tensor rmean, torch::Tensor rvar,
                             torch::Tensor gamma, torch::Tensor beta,
                             double scale, double S, double momentum,
                             double eps, bool training);
"""

_ext = load_inline(
    name="fused_convtranspose3d_bn_globalpool",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["spatial_sums", "eval_finalize", "train_finalize"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    """
    Optimized version of Model. Keeps identical submodules (so state_dicts
    match) but replaces the operator chain with fused custom CUDA kernels.

    Eval mode: the ConvTranspose3d never runs -- its spatial sums are computed
    in closed form from per-channel input sums and kernel sums, then fused
    with scale + BatchNorm + global average pooling in a single tiny kernel.

    Train mode: the convolution runs (cuDNN) and scale + batch statistics +
    running-stat update + BatchNorm + global average pooling are fused into
    one reduction pass plus one tiny finalize kernel (no extra passes over the
    ~130 MB intermediate, which is also never re-read).
    """

    def __init__(self, in_channels, out_channels, kernel_size, scale_factor,
                 eps=1e-5, momentum=0.1):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size)
        self.scale_factor = float(scale_factor)
        self.batch_norm = nn.BatchNorm3d(out_channels, eps=eps, momentum=momentum)
        self.global_avg_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
        ct = self.conv_transpose
        # The closed-form sum decomposition requires the transposed conv to be
        # a "full" convolution (stride 1, no padding, unit dilation/groups).
        self._simple_conv = (
            ct.groups == 1
            and tuple(ct.stride) == (1, 1, 1)
            and tuple(ct.padding) == (0, 0, 0)
            and tuple(ct.dilation) == (1, 1, 1)
            and tuple(ct.output_padding) == (0, 0, 0)
        )

    def _reference_forward(self, x):
        x = self.conv_transpose(x)
        x = x * self.scale_factor
        x = self.batch_norm(x)
        x = self.global_avg_pool(x)
        return x

    def forward(self, x):
        bn = self.batch_norm
        ct = self.conv_transpose

        # Fall back to the reference pipeline for unsupported configurations.
        if not (x.is_cuda and x.dtype == torch.float32):
            return self._reference_forward(x)
        if bn.training and bn.momentum is None:
            return self._reference_forward(x)  # cumulative moving average
        if (not bn.training) and bn.track_running_stats and (
            bn.running_mean is None or not bn.running_mean.is_cuda
        ):
            return self._reference_forward(x)

        B = x.shape[0]
        gamma = bn.weight if bn.affine else x.new_empty(0)
        beta = bn.bias if bn.affine else x.new_empty(0)

        if (not bn.training) and bn.track_running_stats and self._simple_conv:
            # ---- closed-form path: no convolution at all -------------------
            In, Out = ct.in_channels, ct.out_channels
            Sx = _ext.spatial_sums(x, False)[0]            # [B, In]
            Sw = _ext.spatial_sums(ct.weight, False)[0]    # [In, Out]
            D, H, W = x.shape[2], x.shape[3], x.shape[4]
            k = ct.kernel_size
            S_out = (D + k[0] - 1) * (H + k[1] - 1) * (W + k[2] - 1)
            bias = ct.bias if ct.bias is not None else x.new_empty(0)
            out = _ext.eval_finalize(
                Sx, Sw, bias, bn.running_mean, bn.running_var, gamma, beta,
                self.scale_factor, 1.0 / float(S_out), bn.eps,
            )
            return out.view(B, Out, 1, 1, 1)

        # ---- general path: conv + fused reduction --------------------------
        y = ct(x)
        C = y.shape[1]
        S = y.shape[2] * y.shape[3] * y.shape[4]
        sums, sumsq = _ext.spatial_sums(y, True)
        if bn.track_running_stats:
            rmean, rvar = bn.running_mean, bn.running_var
        else:
            rmean = rvar = x.new_empty(0)
        momentum = bn.momentum if bn.momentum is not None else -1.0
        if bn.training and bn.track_running_stats and bn.num_batches_tracked is not None:
            bn.num_batches_tracked += 1
        out = _ext.train_finalize(
            sums, sumsq, rmean, rvar, gamma, beta,
            self.scale_factor, float(S), momentum, bn.eps, bn.training,
        )
        return out.view(B, C, 1, 1, 1)
```

Write an improved complete ModelNew: fix any error, and if the solution is already correct, make it faster. Output the new code in one ```python codeblock. Just the model code, no other text, and NO testing code.
