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 6 of 6 for this task. Your previous solution (turn 5) 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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Optimizations vs. the previous version:
#  1) stats_nhwc_kernel: the double-precision shared-memory atomicAdd CAS
#     loop is replaced by an in-warp shuffle reduction + float shared
#     staging, with only 3 global double atomics per channel per block.
#  2) apply_transpose_nhwc_kernel: rewritten as a 64(s) x 32(c) shared-mem
#     tile where BOTH the global loads (NHWC float4) and the global stores
#     (NCHW, 32 consecutive spatial positions per warp) are fully coalesced
#     128B/warp transactions. The old version issued strided 4B stores
#     (~25% store efficiency) which made it 4.5x slower than the traffic
#     bound. Scalar stores are used because rows of the NCHW output are not
#     16B-aligned when S is not a multiple of 4 (S = 31^3 here).
#  3) Workspace zero-init uses cudaMemsetAsync instead of a fill kernel.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

// ---------------------------------------------------------------------------
// NCHW fallback kernels (layout-generic path)
// ---------------------------------------------------------------------------
__global__ void bn_ms_stats_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,
    double* __restrict__ chan_sum,
    double* __restrict__ chan_sumsq,
    int OC, int slice_elems)
{
    const int slice = blockIdx.y;
    const int oc = slice % OC;
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    float s = 0.f, sq = 0.f;
    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        s  += v.x + v.y + v.z + v.w;
        sq += v.x * v.x + v.y * v.y + v.z * v.z + v.w * v.w;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const float v = y[base + threadIdx.x];
            s += v; sq += v * v;
        }
        if (threadIdx.x < tail) {
            const float v = y[base + head + ((long)vec << 2) + threadIdx.x];
            s += v; sq += v * v;
        }
    }

    __shared__ double bsum[THREADS];
    __shared__ double bsq[THREADS];
    bsum[threadIdx.x] = (double)s;
    bsq[threadIdx.x]  = (double)sq;
    __syncthreads();
    #pragma unroll
    for (int st = THREADS / 2; st > 0; st >>= 1) {
        if (threadIdx.x < st) {
            bsum[threadIdx.x] += bsum[threadIdx.x + st];
            bsq[threadIdx.x]  += bsq[threadIdx.x + st];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        atomicAdd(&slice_sum[slice], bsum[0]);
        atomicAdd(&chan_sum[oc],     bsum[0]);
        atomicAdd(&chan_sumsq[oc],   bsq[0]);
    }
}

__global__ void bn_ms_apply_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 sc = scale[oc];
    const float m  = slice_mean[slice];
    const long base = (long)slice * slice_elems;

    int head = (int)(base & 3L);
    head = head ? (4 - head) : 0;
    if (head > slice_elems) head = slice_elems;
    const int rem  = slice_elems - head;
    const int vec  = rem >> 2;
    const int tail = rem & 3;

    const float4* y4 = reinterpret_cast<const float4*>(y + base + head);
    float4* o4 = reinterpret_cast<float4*>(out + base + head);
    for (int i = blockIdx.x * THREADS + threadIdx.x;
         i < vec; i += gridDim.x * THREADS) {
        const float4 v = y4[i];
        float4 r;
        r.x = sc * (v.x - m);
        r.y = sc * (v.y - m);
        r.z = sc * (v.z - m);
        r.w = sc * (v.w - m);
        o4[i] = r;
    }
    if (blockIdx.x == 0) {
        if (threadIdx.x < head) {
            const long idx = base + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
        if (threadIdx.x < tail) {
            const long idx = base + head + ((long)vec << 2) + threadIdx.x;
            out[idx] = sc * (y[idx] - m);
        }
    }
}

// ---------------------------------------------------------------------------
// Shared finalize kernel: per-slice means + per-channel scales (+running stats)
// ---------------------------------------------------------------------------
__global__ void bn_ms_finalize_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,
    float* __restrict__ scale,
    int NSLICE, int OC,
    double inv_spatial, double inv_total, double total,
    float eps, float momentum,
    int use_batch, int update_running)
{
    const int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < NSLICE) {
        slice_mean[i] = (float)(slice_sum[i] * inv_spatial);
    }
    if (i < OC) {
        float sc;
        if (use_batch) {
            const double mu = chan_sum[i] * inv_total;
            double var = chan_sumsq[i] * inv_total - mu * mu;
            if (var < 0.0) var = 0.0;
            sc = gamma[i] * (float)rsqrt(var + (double)eps);
            if (update_running) {
                running_mean[i] = (1.f - momentum) * running_mean[i] + momentum * (float)mu;
                const double unb = var * (total / (total - 1.0));
                running_var[i]  = (1.f - momentum) * running_var[i] + momentum * (float)unb;
            }
        } else {
            sc = gamma[i] * rsqrtf(running_var[i] + eps);
        }
        scale[i] = sc;
    }
}

// ---------------------------------------------------------------------------
// NHWC (channels_last_3d) kernels. y is logically [N, C, D, H, W] stored as
// [N, S=D*H*W, C] with C == 32 (guarantees 16B alignment of every float4).
// ---------------------------------------------------------------------------

// Pass 1 (NHWC): per-(n,c) spatial sums + per-channel sum / sum-of-squares.
// Block: 256 threads = 32 spatial lanes x 8 float4 channel groups.
// Reduction: in-warp shfl over the 4 spatial lanes sharing a channel group,
// float shared staging across 8 warps, 3 global double atomics per channel.
__global__ void stats_nhwc_kernel(
    const float* __restrict__ y,
    double* __restrict__ slice_sum,   // [N*C]
    double* __restrict__ chan_sum,    // [C]
    double* __restrict__ chan_sumsq,  // [C]
    int C, long S)
{
    const int n = blockIdx.y;
    const long base = (long)n * S * C;
    const int c0 = (threadIdx.x & 7) << 2;   // 0,4,...,28
    const int s_lane = threadIdx.x >> 3;     // 0..31

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float q0 = 0.f, q1 = 0.f, q2 = 0.f, q3 = 0.f;

    const int tiles = (int)((S + 31) >> 5);
    for (int t = blockIdx.x; t < tiles; t += gridDim.x) {
        const long s = ((long)t << 5) + s_lane;
        if (s < S) {
            const float4 v = *reinterpret_cast<const float4*>(y + base + s * C + c0);
            s0 += v.x; s1 += v.y; s2 += v.z; s3 += v.w;
            q0 = fmaf(v.x, v.x, q0);
            q1 = fmaf(v.y, v.y, q1);
            q2 = fmaf(v.z, v.z, q2);
            q3 = fmaf(v.w, v.w, q3);
        }
    }

    // Reduce over the 4 spatial lanes that share a channel group (lanes
    // l, l+8, l+16, l+24 have identical channel group).
    #pragma unroll
    for (int off = 16; off >= 8; off >>= 1) {
        s0 += __shfl_down_sync(0xffffffffu, s0, off);
        s1 += __shfl_down_sync(0xffffffffu, s1, off);
        s2 += __shfl_down_sync(0xffffffffu, s2, off);
        s3 += __shfl_down_sync(0xffffffffu, s3, off);
        q0 += __shfl_down_sync(0xffffffffu, q0, off);
        q1 += __shfl_down_sync(0xffffffffu, q1, off);
        q2 += __shfl_down_sync(0xffffffffu, q2, off);
        q3 += __shfl_down_sync(0xffffffffu, q3, off);
    }

    __shared__ float sh_sum[8][32];
    __shared__ float sh_sq[8][32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane < 8) {
        sh_sum[warp][c0 + 0] = s0; sh_sum[warp][c0 + 1] = s1;
        sh_sum[warp][c0 + 2] = s2; sh_sum[warp][c0 + 3] = s3;
        sh_sq[warp][c0 + 0]  = q0; sh_sq[warp][c0 + 1]  = q1;
        sh_sq[warp][c0 + 2]  = q2; sh_sq[warp][c0 + 3]  = q3;
    }
    __syncthreads();

    if (threadIdx.x < 32) {
        const int c = threadIdx.x;
        float fs = 0.f, fq = 0.f;
        #pragma unroll
        for (int w = 0; w < 8; ++w) {
            fs += sh_sum[w][c];
            fq += sh_sq[w][c];
        }
        const double ds = (double)fs;
        atomicAdd(&slice_sum[n * C + c], ds);
        atomicAdd(&chan_sum[c],          ds);
        atomicAdd(&chan_sumsq[c],        (double)fq);
    }
}

// Pass 2 (NHWC -> NCHW transpose with fused affine):
//   out[n, c, s] = scale[c] * (y[n, s, c] - slice_mean[n, c])
// 64(s) x 32(c) shared-memory tile. Loads: coalesced float4 over channels.
// Stores: each warp writes 32 consecutive spatial positions of one channel
// (128B/warp, fully coalesced, no alignment requirement). All shared-mem
// accesses are bank-conflict free.
__global__ void apply_transpose_nhwc_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const float* __restrict__ scale,      // [C]
    const float* __restrict__ slice_mean, // [N*C]
    int C, long S)
{
    __shared__ float tile[32][65];
    const int n = blockIdx.y;
    const long s0 = (long)blockIdx.x << 6;

    const int cg = threadIdx.x & 7;        // float4 channel group 0..7
    const int s_sub = threadIdx.x >> 3;    // 0..31
    const float4 sc4 = *reinterpret_cast<const float4*>(scale + (cg << 2));
    const float4 mn4 = *reinterpret_cast<const float4*>(slice_mean + n * C + (cg << 2));
    const long nbase = (long)n * S;

    #pragma unroll
    for (int half = 0; half < 2; ++half) {
        const int ss = s_sub + (half << 5);   // 0..63
        const long s = s0 + ss;
        if (s < S) {
            const float4 v = __ldcs(reinterpret_cast<const float4*>(
                y + (nbase + s) * C + (cg << 2)));
            tile[(cg << 2) + 0][ss] = sc4.x * (v.x - mn4.x);
            tile[(cg << 2) + 1][ss] = sc4.y * (v.y - mn4.y);
            tile[(cg << 2) + 2][ss] = sc4.z * (v.z - mn4.z);
            tile[(cg << 2) + 3][ss] = sc4.w * (v.w - mn4.w);
        }
    }
    __syncthreads();

    const long rem = S - s0;
    const int s = threadIdx.x & 63;
    if (s < rem) {
        int c = threadIdx.x >> 6;   // 0..3, strides of 4 over 32 channels
        float* obase = out + ((long)n * C) * S + s0 + s;
        #pragma unroll
        for (int i = 0; i < 8; ++i, c += 4) {
            __stcs(obase + (long)c * S, tile[c][s]);
        }
    }
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
static inline void launch_finalize(
    double* slice_sum_p, double* chan_sum_p, double* chan_sumsq_p,
    float* gamma_p, float* rm_p, float* rv_p,
    float* slice_mean_p, float* scale_p,
    int NSLICE, int OC, int slice_elems, int N,
    float eps, float momentum, bool use_batch, bool update_running,
    cudaStream_t stream)
{
    const double spatial = (double)slice_elems;
    const double total   = spatial * (double)N;
    const int work = NSLICE > OC ? NSLICE : OC;
    bn_ms_finalize_kernel<<<(work + 255) / 256, 256, 0, stream>>>(
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        gamma_p, rm_p, rv_p,
        slice_mean_p, scale_p,
        NSLICE, OC, 1.0 / spatial, 1.0 / total, total,
        eps, momentum,
        use_batch ? 1 : 0, update_running ? 1 : 0);
}

// NCHW path (fallback)
torch::Tensor fused_bn_meansub(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub expects CUDA float32 input");
    auto yc = y.contiguous();
    const int N  = yc.size(0);
    const int OC = yc.size(1);
    const long slice_elems_l = yc.numel() / ((long)N * (long)OC);
    const int slice_elems = (int)slice_elems_l;
    const int NSLICE = N * OC;

    auto out = torch::empty_like(yc);
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, yc.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(yc.device());
    auto dbuf = torch::empty({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, yc.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    const int max_vec = (slice_elems + 3) / 4;
    int chunks = (max_vec + THREADS * 4 - 1) / (THREADS * 4);
    if (chunks < 1) chunks = 1;
    if (chunks > 65535) chunks = 65535;
    dim3 grid(chunks, NSLICE);

    bn_ms_stats_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, slice_elems);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, slice_elems, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    bn_ms_apply_kernel<<<grid, THREADS, 0, stream>>>(
        yc.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, slice_elems);

    return out;
}

// NHWC (channels_last_3d) fast path; requires C == 32.
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y,
    torch::Tensor gamma,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine)
{
    TORCH_CHECK(y.is_cuda() && y.dtype() == torch::kFloat32,
                "fused_bn_meansub_cl expects CUDA float32 input");
    TORCH_CHECK(y.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "fused_bn_meansub_cl expects channels_last_3d input");
    const int N  = y.size(0);
    const int OC = y.size(1);
    TORCH_CHECK(OC == 32, "fused_bn_meansub_cl requires 32 channels");
    const long S = y.numel() / ((long)N * (long)OC);
    const int NSLICE = N * OC;

    // final output: NCHW contiguous
    auto out = torch::empty({N, OC, y.size(2), y.size(3), y.size(4)}, y.options());
    auto gc = has_affine ? gamma.contiguous()
                         : torch::ones({OC}, y.options());

    auto opts_d = torch::TensorOptions().dtype(torch::kFloat64).device(y.device());
    auto dbuf = torch::empty({(long)NSLICE + 2L * OC}, opts_d);
    auto fbuf = torch::empty({(long)NSLICE + OC}, y.options());
    double* slice_sum_p  = dbuf.data_ptr<double>();
    double* chan_sum_p   = slice_sum_p + NSLICE;
    double* chan_sumsq_p = chan_sum_p + OC;
    float*  slice_mean_p = fbuf.data_ptr<float>();
    float*  scale_p      = slice_mean_p + NSLICE;

    auto stream = at::cuda::getCurrentCUDAStream();
    cudaMemsetAsync(dbuf.data_ptr(), 0, dbuf.numel() * sizeof(double), stream.stream());

    const int tiles32 = (int)((S + 31) >> 5);
    int gx = tiles32 < 32 ? tiles32 : 32;
    dim3 sgrid(gx, N);
    stats_nhwc_kernel<<<sgrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(),
        slice_sum_p, chan_sum_p, chan_sumsq_p,
        OC, S);

    launch_finalize(slice_sum_p, chan_sum_p, chan_sumsq_p,
        gc.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        slice_mean_p, scale_p,
        NSLICE, OC, (int)S, N,
        (float)eps, (float)momentum, use_batch, update_running,
        stream.stream());

    dim3 agrid((unsigned)((S + 63) >> 6), N);
    apply_transpose_nhwc_kernel<<<agrid, THREADS, 0, stream>>>(
        y.data_ptr<float>(), out.data_ptr<float>(),
        scale_p, slice_mean_p,
        OC, S);

    return out;
}
"""

cpp_source = """
torch::Tensor fused_bn_meansub(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
torch::Tensor fused_bn_meansub_cl(
    torch::Tensor y, torch::Tensor gamma,
    torch::Tensor running_mean, torch::Tensor running_var,
    double eps, double momentum,
    bool use_batch, bool update_running, bool has_affine);
"""

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


class ModelNew(nn.Module):
    """
    ConvTranspose3d (cudnn, channels_last_3d, no bias kernel needed) followed
    by a fused custom-CUDA BatchNorm3d + spatial-mean-subtraction tail.

    Exact algebra used:
      BN(y)[n,c,s] - mean_s BN(y) == scale[c] * (g[n,c,s] - mean_s g),
    where y = g + bias[c], scale[c] = gamma[c] / sqrt(var_c + eps). The
    per-channel conv bias, BN beta and BN channel mean all cancel exactly,
    so the conv is run with bias=None and the NHWC conv output is consumed
    directly by the fused kernels (stats + transposing apply).
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=True):
        super(ModelNew, self).__init__()
        torch.backends.cudnn.benchmark = True
        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):
        bn = self.batch_norm
        ct = self.conv_transpose

        if (x.is_cuda and x.dtype == torch.float32
                and bn.momentum is not None and ct.groups == 1):
            track = bn.track_running_stats
            use_batch = bool(bn.training) or (not track)
            update_running = bool(bn.training) and bool(track)

            if track:
                rm = bn.running_mean
                rv = bn.running_var
            else:
                rm = torch.zeros(bn.num_features, device=x.device, dtype=torch.float32)
                rv = torch.ones(bn.num_features, device=x.device, dtype=torch.float32)

            if update_running and bn.num_batches_tracked is not None:
                bn.num_batches_tracked += 1

            gamma = bn.weight if bn.affine else torch.empty(0, device=x.device)
            eps = float(bn.eps)
            mom = float(bn.momentum)

            # Bias cancels exactly in BN->mean-sub, so run the conv without it.
            x_cl = x.contiguous(memory_format=torch.channels_last_3d)
            w_cl = ct.weight.contiguous(memory_format=torch.channels_last_3d)
            y = F.conv_transpose3d(
                x_cl, w_cl, None, ct.stride, ct.padding,
                ct.output_padding, ct.groups, ct.dilation)

            if (y.is_contiguous(memory_format=torch.channels_last_3d)
                    and ct.out_channels == 32):
                return fused_op.fused_bn_meansub_cl(
                    y, gamma, rm, rv, eps, mom,
                    use_batch, update_running, bool(bn.affine))

            # layout fallback: NCHW fused kernels
            return fused_op.fused_bn_meansub(
                y.contiguous(), gamma, rm, rv, eps, mom,
                use_batch, update_running, bool(bn.affine))

        # pure-PyTorch fallback (keeps full bias/BN semantics)
        y = self.conv_transpose(x)
        z = self.batch_norm(y)
        return z - torch.mean(z, dim=(2, 3, 4), keepdim=True)
```
Result of evaluating it:
The solution is correct on the test inputs. Mean runtime 1.56 ms; the PyTorch reference (with TF32 enabled) takes 4.19 ms; speedup 2.686x.
Profiler, top device kernels of one forward pass of your solution:
| device kernel | total device time (us) | calls |
|---|---|---|
| sm80_xmma_dgrad_implicit_gemm_indexed_tf32f32_tf32f32_f32_nhwckrsc_nhwc_tilesize256x64x32_ | 740.4 | 1 |
| apply_transpose_nhwc_kernel(float const*, float*, float const*, float const*, int, long) | 381.9 | 1 |
| stats_nhwc_kernel(float const*, double*, double*, double*, int, long) | 199.2 | 1 |
| void xmma__5x_cudnn::implicit_gemm::strided_dgrad_indexed::kernel_helper_stage_3<xmma__5x_ | 83.1 | 1 |
| void at::native::elementwise_kernel<128, 2, at::native::gpu_kernel_impl_nocast<at::native: | 41.0 | 2 |
| void cudnn::engines_precompiled::convertTensor_kernel<float, float, float, (cudnnKernelDat | 27.8 | 2 |
| void xmma__5x_cudnn::implicit_gemm::strided_dgrad_indexed::kernel_helper_stage_2<xmma__5x_ | 4.7 | 1 |
| bn_ms_finalize_kernel(double const*, double const*, double const*, float const*, float*, f | 4.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.
