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):
    """
    A 3D convolutional transpose layer followed by Batch Normalization and subtraction.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)
        self.batch_norm = nn.BatchNorm3d(out_channels)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = x - torch.mean(x, dim=(2, 3, 4), keepdim=True)  # Subtract mean along spatial dimensions
        return x

batch_size = 16
in_channels = 16
out_channels = 32
depth, height, width = 16, 32, 32
kernel_size = 3
stride = 2
padding = 1

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

def get_init_inputs():
    return [in_channels, out_channels, kernel_size, stride, padding]
    ```
    
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 2 of 6 for this task. Your previous solution (turn 1) 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

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

// ---------------------------------------------------------------------------
// Kernel A: fused ConvTranspose3d (k=3, s=2, p=1) + statistics accumulation.
// Each thread computes a 2(d) x 2(h) x 4(w) tile of outputs (16 outputs).
// Block-level partial sums -> atomics for:
//   - per (n,c) spatial sum (for the mean subtraction)
//   - per channel sum / sumsq (for batch statistics in training mode)
// ---------------------------------------------------------------------------
__global__ void convt3d_k3s2p1_stats_kernel(
    const float* __restrict__ x,        // [N, IC, ID, IH, IW]
    const float* __restrict__ wpad,     // [IC, OC, 9, 4] (kw padded to 4)
    const float* __restrict__ bias,     // [OC]
    float* __restrict__ y,              // [N, OC, OD, OH, OW]
    double* __restrict__ slice_sum,     // [N*OC]
    double* __restrict__ chan_sum,      // [OC]
    double* __restrict__ chan_sumsq,    // [OC]
    int IC, int OC,
    int ID, int IH, int IW,
    int OD, int OH, int OW,
    int tiles_d, int tiles_h, int tiles_w)
{
    const int slice = blockIdx.y;           // n*OC + oc
    const int n  = slice / OC;
    const int oc = slice - n * OC;

    const int tiles_per_slice = tiles_d * tiles_h * tiles_w;
    const int tile_idx = blockIdx.x * blockDim.x + threadIdx.x;

    float lsum = 0.f;
    float lsq  = 0.f;

    if (tile_idx < tiles_per_slice) {
        const int tw_i = tile_idx % tiles_w;
        const int th_i = (tile_idx / tiles_w) % tiles_h;
        const int td_i = tile_idx / (tiles_w * tiles_h);

        const int id_b = td_i;
        const int ih_b = th_i;
        const int iw_b = 2 * tw_i;

        const int od0 = 2 * td_i;
        const int oh0 = 2 * th_i;
        const int ow0 = 4 * tw_i;

        const float* xn = x + (long)n * IC * ID * IH * IW;
        const long xICs = (long)ID * IH * IW;

        float acc[2][2][4];
        #pragma unroll
        for (int a = 0; a < 2; ++a)
            #pragma unroll
            for (int b = 0; b < 2; ++b)
                #pragma unroll
                for (int c = 0; c < 4; ++c) acc[a][b][c] = 0.f;

        for (int ic = 0; ic < IC; ++ic) {
            const float* xic = xn + (long)ic * xICs;
            const float* wic = wpad + ((long)ic * OC + oc) * 36;

            // load 2x2x3 input tile into registers (zero-filled out of range)
            float xv[2][2][3];
            #pragma unroll
            for (int ad = 0; ad < 2; ++ad) {
                const int id = id_b + ad;
                const int idc = min(id, ID - 1);
                const bool okd = (id < ID);
                #pragma unroll
                for (int ah = 0; ah < 2; ++ah) {
                    const int ih = ih_b + ah;
                    const int ihc = min(ih, IH - 1);
                    const bool okh = okd && (ih < IH);
                    const float* xr = xic + ((long)idc * IH + ihc) * IW;
                    #pragma unroll
                    for (int aw = 0; aw < 3; ++aw) {
                        const int iw = iw_b + aw;
                        const bool okw = okh && (iw < IW);
                        xv[ad][ah][aw] = okw ? xr[min(iw, IW - 1)] : 0.f;
                    }
                }
            }

            // 27 weights (as 9 float4) applied to the tile
            #pragma unroll
            for (int kd = 0; kd < 3; ++kd) {
                const int dm = (kd == 1) ? 0 : 1;
                const int ai = (kd == 0) ? 1 : 0;
                #pragma unroll
                for (int kh = 0; kh < 3; ++kh) {
                    const int hm = (kh == 1) ? 0 : 1;
                    const int hi = (kh == 0) ? 1 : 0;
                    const float4 wv = *reinterpret_cast<const float4*>(wic + (kd * 3 + kh) * 4);
                    acc[dm][hm][0] += xv[ai][hi][0] * wv.y;   // kw = 1
                    acc[dm][hm][2] += xv[ai][hi][1] * wv.y;   // kw = 1
                    acc[dm][hm][1] += xv[ai][hi][0] * wv.z;   // kw = 2
                    acc[dm][hm][3] += xv[ai][hi][1] * wv.z;   // kw = 2
                    acc[dm][hm][1] += xv[ai][hi][1] * wv.x;   // kw = 0
                    acc[dm][hm][3] += xv[ai][hi][2] * wv.x;   // kw = 0
                }
            }
        }

        const float bv = bias[oc];
        const long ybase = (long)slice * OD * OH * OW;
        #pragma unroll
        for (int a = 0; a < 2; ++a) {
            const int od = od0 + a;
            if (od >= OD) continue;
            #pragma unroll
            for (int b = 0; b < 2; ++b) {
                const int oh = oh0 + b;
                if (oh >= OH) continue;
                #pragma unroll
                for (int c = 0; c < 4; ++c) {
                    const int ow = ow0 + c;
                    if (ow >= OW) continue;
                    const float v = acc[a][b][c] + bv;
                    y[ybase + ((long)od * OH + oh) * OW + ow] = v;
                    lsum += v;
                    lsq  += v * v;
                }
            }
        }
    }

    // block reduction of the two partial statistics
    __shared__ double s_sum[256];
    __shared__ double s_sq[256];
    const int tid = threadIdx.x;
    s_sum[tid] = (double)lsum;
    s_sq[tid]  = (double)lsq;
    __syncthreads();
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            s_sum[tid] += s_sum[tid + s];
            s_sq[tid]  += s_sq[tid + s];
        }
        __syncthreads();
    }
    if (tid == 0) {
        atomicAdd(&slice_sum[slice], s_sum[0]);
        atomicAdd(&chan_sum[oc],     s_sum[0]);
        atomicAdd(&chan_sumsq[oc],   s_sq[0]);
    }
}

// ---------------------------------------------------------------------------
// Kernel B: finalize statistics -> per-slice means, per-channel scales,
// and (optionally) update running statistics.
// ---------------------------------------------------------------------------
__global__ void finalize_stats_kernel(
    const double* __restrict__ slice_sum,
    const double* __restrict__ chan_sum,
    const double* __restrict__ chan_sumsq,
    const float* __restrict__ gamma,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ slice_mean,  // [N*OC]
    float* __restrict__ scale,       // [OC]
    int NSLICE, int OC,
    double inv_spatial, double inv_total, double total,
    float eps, float momentum,
    int training, int update_running)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < NSLICE) {
        slice_mean[i] = (float)((double)slice_sum[i] * inv_spatial);
    }
    if (i < OC) {
        float s;
        if (training) {
            const double mu = (double)chan_sum[i] * inv_total;
            double var = (double)chan_sumsq[i] * inv_total - mu * mu;
            if (var < 0.0) var = 0.0;
            s = gamma[i] * (float)rsqrt(var + (double)eps);
            if (update_running) {
                const float rm = running_mean[i];
                running_mean[i] = (1.f - momentum) * rm + momentum * (float)mu;
                const double unb = var * (total / (total - 1.0));
                const float rv = running_var[i];
                running_var[i] = (1.f - momentum) * rv + momentum * (float)unb;
            }
        } else {
            s = gamma[i] * rsqrtf(running_var[i] + eps);
        }
        scale[i] = s;
    }
}

// ---------------------------------------------------------------------------
// Kernel C: out = scale[c] * (y - slice_mean[n,c])  (fused BN + mean subtract)
// ---------------------------------------------------------------------------
__global__ void apply_scale_mean_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,
    const float* __restrict__ slice_mean,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const float s = scale[oc];
    const float m = slice_mean[slice];
    const long base = (long)slice * slice_elems;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < slice_elems; i += blockDim.x * gridDim.x) {
        out[base + i] = s * (y[base + i] - m);
    }
}

// ---------------------------------------------------------------------------
// Host wrapper
// ---------------------------------------------------------------------------
torch::Tensor fused_convt_bn_meansub(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor bias,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool training, bool track_running, bool has_bias, bool has_affine)
{
    TORCH_CHECK(x.is_cuda() && x.dtype() == torch::kFloat32);
    auto xc = x.contiguous();
    auto wc = weight.contiguous();
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({weight.size(1)}, x.options());

    const int N  = xc.size(0);
    const int IC = xc.size(1);
    const int ID = xc.size(2);
    const int IH = xc.size(3);
    const int IW = xc.size(4);
    const int OC = wc.size(1);

    const int OD = (ID - 1) * 2 - 2 + 3;
    const int OH = (IH - 1) * 2 - 2 + 3;
    const int OW = (IW - 1) * 2 - 2 + 3;

    auto y   = torch::empty({N, OC, OD, OH, OW}, x.options());
    auto out = torch::empty_like(y);

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto slice_sum  = torch::zeros({(long)N * OC}, opts_d);
    auto chan_sum   = torch::zeros({OC}, opts_d);
    auto chan_sumsq = torch::zeros({OC}, opts_d);
    auto slice_mean = torch::empty({(long)N * OC}, x.options());
    auto scale      = torch::empty({OC}, x.options());

    auto wpad = torch::constant_pad_nd(wc.view({IC, OC, 9, 3}), {0, 1}, 0.0)
                    .contiguous();

    auto bc = has_bias ? bias.contiguous()
                       : torch::zeros({OC}, x.options());

    auto stream = at::cuda::getCurrentCUDAStream();

    // Kernel A
    const int tiles_d = (OD + 1) / 2;
    const int tiles_h = (OH + 1) / 2;
    const int tiles_w = (OW + 3) / 4;
    const int tiles_per_slice = tiles_d * tiles_h * tiles_w;
    const int threads = 256;
    const int blocks_x = (tiles_per_slice + threads - 1) / threads;
    dim3 gridA(blocks_x, N * OC);
    convt3d_k3s2p1_stats_kernel<<<gridA, threads, 0, stream>>>(
        xc.data_ptr<float>(), wpad.data_ptr<float>(), bc.data_ptr<float>(),
        y.data_ptr<float>(),
        slice_sum.data_ptr<double>(), chan_sum.data_ptr<double>(),
        chan_sumsq.data_ptr<double>(),
        IC, OC, ID, IH, IW, OD, OH, OW,
        tiles_d, tiles_h, tiles_w);

    // Kernel B
    const int NSLICE = N * OC;
    const double spatial = (double)OD * OH * OW;
    const double total = spatial * (double)N;
    const int work = NSLICE > OC ? NSLICE : OC;
    const int blocksB = (work + 255) / 256;
    const int update_running = (training && track_running) ? 1 : 0;
    finalize_stats_kernel<<<blocksB, 256, 0, stream>>>(
        slice_sum.data_ptr<double>(), chan_sum.data_ptr<double>(),
        chan_sumsq.data_ptr<double>(),
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean.data_ptr<float>(), scale.data_ptr<float>(),
        NSLICE, OC, 1.0 / spatial, 1.0 / total, total,
        (float)eps, (float)momentum, training ? 1 : 0, update_running);

    // Kernel C
    const int slice_elems = OD * OH * OW;
    int cx = (slice_elems + threads - 1) / threads;
    if (cx > 65535) cx = 65535;
    dim3 gridC(cx, NSLICE);
    apply_scale_mean_kernel<<<gridC, threads, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(),
        scale.data_ptr<float>(), slice_mean.data_ptr<float>(),
        OC, slice_elems);

    return out;
}
"""

cpp_source = """
torch::Tensor fused_convt_bn_meansub(
    torch::Tensor x, torch::Tensor weight, torch::Tensor bias,
    torch::Tensor gamma, torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool training, bool track_running, bool has_bias, bool has_affine);
"""

fused_op = load_inline(
    name="fused_convt_bn_meansub",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_convt_bn_meansub"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    ConvTranspose3d + BatchNorm3d + spatial-mean subtraction, fused into
    custom CUDA kernels. Falls back to native PyTorch ops for unsupported
    configurations.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size,
            stride=stride, padding=padding, bias=bias)
        self.batch_norm = nn.BatchNorm3d(out_channels)

        ct = self.conv_transpose
        self._use_fused = (
            ct.kernel_size == (3, 3, 3)
            and ct.stride == (2, 2, 2)
            and ct.padding == (1, 1, 1)
            and ct.dilation == (1, 1, 1)
            and ct.groups == 1
            and ct.output_padding == (0, 0, 0)
            and ct.padding_mode == "zeros"
        )

    def forward(self, x):
        bn = self.batch_norm
        if (self._use_fused and x.is_cuda and x.dtype == torch.float32):
            ct = self.conv_transpose
            track = bn.track_running_stats
            if not track and not bn.training:
                # no running stats available in eval mode -> native fallback
                pass
            else:
                rm = bn.running_mean if track else torch.zeros(
                    bn.num_features, device=x.device)
                rv = bn.running_var if track else torch.ones(
                    bn.num_features, device=x.device)
                return fused_op.fused_convt_bn_meansub(
                    x, ct.weight,
                    ct.bias if ct.bias is not None else torch.empty(0, device=x.device),
                    bn.weight if bn.affine else torch.empty(0, device=x.device),
                    rm, rv,
                    float(bn.eps), float(bn.momentum),
                    bool(bn.training), bool(track),
                    ct.bias is not None, bool(bn.affine))

        # fallback: original PyTorch path
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = x - torch.mean(x, dim=(2, 3, 4), keepdim=True)
        return x
```
Result of evaluating it:
The solution compiled and ran but its output did not match the reference. Details:
```
{"avg_difference": ["0.000227", "0.000227", "0.000227", "0.000227", "0.000227"], "correctness_issue": "Output mismatch", "correctness_trials": "(0 / 5)", "max_difference": ["0.003292", "0.003529", "0.003744", "0.003445", "0.003541"]}
```

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.
