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 5 of 6 for this task. Your previous solution (turn 4) 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) The transposed conv is run in channels_last_3d (NHWC), which is the layout
#     cudnn's dgrad engine uses natively. This eliminates the big NHWC->NCHW
#     transpose kernel (~396us) and the input transforms.
#  2) The conv bias-add kernel (~378us) is removed entirely: since
#     BN(y) - mean_spatial(BN(y)) == scale[c] * (g - mean_spatial(g)) with
#     g = conv_transpose_nobias(x) (a per-channel constant cancels exactly),
#     the bias never has to be materialized. The bias parameter is kept for
#     state_dict compatibility and still used on the pure-PyTorch fallback path.
#  3) New fused tail that consumes NHWC directly: a coalesced stats kernel
#     (per-(n,c) spatial means + per-channel BN statistics) and a tiled
#     shared-memory transpose kernel that applies scale*(g-mean) while writing
#     the final NCHW result.
#  4) The previous NCHW fused kernels are kept as a fallback for any layout.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256

// ---------------------------------------------------------------------------
// NCHW fallback kernels (from previous version, with alignment-safe indexing)
// ---------------------------------------------------------------------------
__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] but stored
// as [N, S=D*H*W, C] with C == 32. Requires 16B alignment (C multiple of 4).
// ---------------------------------------------------------------------------

// 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.
__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);
        }
    }

    __shared__ double sh_sum[32];
    __shared__ double sh_sq[32];
    if (threadIdx.x < 32) {
        sh_sum[threadIdx.x] = 0.0;
        sh_sq[threadIdx.x]  = 0.0;
    }
    __syncthreads();
    atomicAdd(&sh_sum[c0 + 0], (double)s0);
    atomicAdd(&sh_sum[c0 + 1], (double)s1);
    atomicAdd(&sh_sum[c0 + 2], (double)s2);
    atomicAdd(&sh_sum[c0 + 3], (double)s3);
    atomicAdd(&sh_sq[c0 + 0], (double)q0);
    atomicAdd(&sh_sq[c0 + 1], (double)q1);
    atomicAdd(&sh_sq[c0 + 2], (double)q2);
    atomicAdd(&sh_sq[c0 + 3], (double)q3);
    __syncthreads();

    if (threadIdx.x < 32) {
        const int c = threadIdx.x;
        atomicAdd(&slice_sum[n * C + c], sh_sum[c]);
        atomicAdd(&chan_sum[c],          sh_sum[c]);
        atomicAdd(&chan_sumsq[c],        sh_sq[c]);
    }
}

// Pass 2 (NHWC -> NCHW transpose with fused affine):
//   out[n, c, s] = scale[c] * (y[n, s, c] - slice_mean[n, c])
// 32x32 shared-memory tile, conflict-free in both phases.
__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][33];
    const int n = blockIdx.y;
    const long s0 = ((long)blockIdx.x << 5);

    const int c0 = (threadIdx.x & 7) << 2;
    const int s_lane = threadIdx.x >> 3;
    const long s = s0 + s_lane;
    if (s < S) {
        const float4 v  = *reinterpret_cast<const float4*>(y + ((long)n * S + s) * C + c0);
        const float4 sc = *reinterpret_cast<const float4*>(scale + c0);
        const float4 mn = *reinterpret_cast<const float4*>(slice_mean + n * C + c0);
        tile[c0 + 0][s_lane] = sc.x * (v.x - mn.x);
        tile[c0 + 1][s_lane] = sc.y * (v.y - mn.y);
        tile[c0 + 2][s_lane] = sc.z * (v.z - mn.z);
        tile[c0 + 3][s_lane] = sc.w * (v.w - mn.w);
    }
    __syncthreads();

    const int c  = threadIdx.x >> 3;         // 0..31
    const int s4 = (threadIdx.x & 7) << 2;   // 0,4,...,28
    float* orow = out + ((long)n * C + c) * S + s0;
    const long rem = S - s0;
    if (rem >= 32) {
        orow[s4 + 0] = tile[c][s4 + 0];
        orow[s4 + 1] = tile[c][s4 + 1];
        orow[s4 + 2] = tile[c][s4 + 2];
        orow[s4 + 3] = tile[c][s4 + 3];
    } else {
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            if (s4 + i < rem) orow[s4 + i] = tile[c][s4 + i];
        }
    }
}

// ---------------------------------------------------------------------------
// 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::zeros({(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();

    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::zeros({(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();

    const int tiles = (int)((S + 31) >> 5);
    int gx = tiles < 64 ? tiles : 64;
    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(tiles, 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_v3",
    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.82 ms; the PyTorch reference (with TF32 enabled) takes 4.19 ms; speedup 2.302x.
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_ | 741.1 | 1 |
| apply_transpose_nhwc_kernel(float const*, float*, float const*, float const*, int, long) | 631.6 | 1 |
| stats_nhwc_kernel(float const*, double*, double*, double*, int, long) | 216.4 | 1 |
| void xmma__5x_cudnn::implicit_gemm::strided_dgrad_indexed::kernel_helper_stage_3<xmma__5x_ | 83.5 | 1 |
| void at::native::elementwise_kernel<128, 2, at::native::gpu_kernel_impl_nocast<at::native: | 40.6 | 2 |
| void cudnn::engines_precompiled::convertTensor_kernel<float, float, float, (cudnnKernelDat | 27.3 | 2 |
| void xmma__5x_cudnn::implicit_gemm::strided_dgrad_indexed::kernel_helper_stage_2<xmma__5x_ | 4.1 | 1 |
| void xmma__5x_cudnn::implicit_gemm::strided_dgrad_indexed::kernel_helper_stage_1<xmma__5x_ | 4.0 | 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.
