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 model that performs a 3D transposed convolution, followed by batch normalization, 
    two average pooling layers.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x


batch_size = 64
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 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, bias_shape]
    ```
    
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 4 of 6 for this task. Your previous solution (turn 3) was:
```python
(no code was extracted)
```
Result of evaluating it:
Your reply contained no usable ```python code block (the reply was cut off at the length limit).

The fastest correct solution so far (22.547x 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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline


# ============================================================================
# Custom CUDA kernels
#
# Optimizations vs. the previous version:
#  1. Training mode no longer runs a separate full-volume statistics kernel.
#     A single fused kernel computes the raw ConvTranspose3d output for the
#     whole 63^3 volume (organized as 8x4x4 voxel tiles), writes the 4x4x4
#     block sums (pooled pre-BN values) for the interior 60^3 region, and
#     accumulates per-channel sum / sum-of-squares for BatchNorm on the fly.
#  2. Threads are assigned to voxels sorted by parity class so that warps are
#     (nearly) uniform in the number of kernel taps they execute.  This
#     removes the divergent-loop serialization of the naive mapping and
#     roughly doubles the conv throughput.
#  3. BatchNorm scale/shift are finalized in a tiny dedicated kernel that
#     also updates the running statistics, and BN+pool scaling is applied
#     with a single addcmul, so training mode adds only ~50us of overhead.
#  4. Inference mode keeps everything in one kernel: ConvTranspose3d +
#     BatchNorm(running stats) + AvgPool(4,4,4), never materializing the
#     63^3 intermediate in DRAM.
# ============================================================================

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

#define CI 3
#define CO 16
#define IDIM 32
#define ODIM 15
#define NTHREADS 128
#define PITCH (CO + 1)
#define SXD 45   // 5*3*3 input patch per channel for an 8x4x4 output tile

// Parity classes ordered by tap count so warps are near-uniform:
// 0:eee(1 tap) 1:eeo(2) 2:eoe(2) 3:oee(2) 4:eoo(4) 5:oeo(4) 6:ooe(4) 7:ooo(8)
__constant__ int c_pd[8] = {0, 0, 0, 1, 0, 1, 1, 1};
__constant__ int c_ph[8] = {0, 0, 1, 0, 1, 0, 1, 1};
__constant__ int c_pw[8] = {0, 1, 0, 0, 1, 1, 0, 1};

// For local coordinate l (parity dependent):
//   global D = base + l ; input d = (D + 1 - kd) / 2
//   even l: 1 tap  -> a = l/2   , kd = 1
//   odd  l: 2 taps -> a = l/2   , kd = 2
//                     a = l/2+1 , kd = 0
__device__ __forceinline__ int setup_taps(int l, int* a, int* k) {
    if ((l & 1) == 0) { a[0] = l >> 1; k[0] = 1; return 1; }
    a[0] = l >> 1; k[0] = 2; a[1] = (l >> 1) + 1; k[1] = 0;
    return 2;
}

__device__ __forceinline__ void conv_acc(
    const float* s_x, const float* s_w,
    const int* ad, const int* ah, const int* aw,
    const int* kd, const int* kh, const int* kw,
    int nd, int nh, int nw, float* acc)
{
    for (int i = 0; i < nd; ++i) {
        for (int j = 0; j < nh; ++j) {
            for (int m = 0; m < nw; ++m) {
                const int xs = ad[i] * 9 + ah[j] * 3 + aw[m];
                const float x0 = s_x[xs];
                const float x1 = s_x[xs + SXD];
                const float x2 = s_x[xs + 2 * SXD];
                const int woff = kd[i] * 9 + kh[j] * 3 + kw[m];
                #pragma unroll
                for (int oc = 0; oc < CO; ++oc) {
                    const float* wp = s_w + oc * 27 + woff;
                    acc[oc] += x0 * wp[0] + x1 * wp[CO * 27] + x2 * wp[2 * CO * 27];
                }
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Training kernel: raw ConvTranspose3d over the full 63^3 volume, fused
// 4x4x4 block pooling (raw block sums) + BatchNorm statistics accumulation.
// Tiles: 8x4x4 voxels, grid per image = 8*16*16 tiles (covers 64^3 >= 63^3).
// ---------------------------------------------------------------------------
__global__ void fused_train_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    float* __restrict__ out,
    double* __restrict__ gsum,
    double* __restrict__ gsq,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_bsum[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) s_bsum[tid] = 0.0f;

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);

    float tsq[CO];
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) tsq[oc] = 0.0f;
    __syncthreads();

    const int TPI = 8 * 16 * 16;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem >> 8;
        rem          &= 255;
        const int th  = rem >> 4;
        const int tw  = rem & 15;
        const bool interior = (td < 7) && (th < 15) && (tw < 15);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (interior) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a, gj = 2 * th + b, gk = 2 * tw + c;
                float v = 0.0f;
                if (gi < IDIM && gj < IDIM && gk < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + gj) * IDIM + gk];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        const bool valid = interior ||
            ((8 * td + vd) < 63 && (4 * th + vh) < 63 && (4 * tw + vw) < 63);
        if (valid) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) {
            s_red[vid][oc] = acc[oc];
            tsq[oc] += acc[oc] * acc[oc];
        }
        __syncthreads();

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            float s = 0.0f;
            #pragma unroll 8
            for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
            atomicAdd(&s_bsum[oc], s);
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM && th < ODIM && tw < ODIM) {
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s;
            }
        }
    }

    __syncthreads();
    #pragma unroll
    for (int oc = 0; oc < CO; ++oc) s_red[tid][oc] = tsq[oc];
    __syncthreads();
    if (tid < CO) {
        float q = 0.0f;
        #pragma unroll 8
        for (int i = 0; i < NTHREADS; ++i) q += s_red[i][tid];
        atomicAdd(gsq + tid, (double)q);
        atomicAdd(gsum + tid, (double)s_bsum[tid]);
    }
}

// ---------------------------------------------------------------------------
// Inference kernel: ConvTranspose3d + BatchNorm(scale/shift) + AvgPool(4)
// fully fused.  Only voxels in the 60^3 region that pooling reads are
// computed.  Tiles per image: 8*15*15.
// ---------------------------------------------------------------------------
__global__ void fused_eval_kernel(
    const float* __restrict__ x,
    const float* __restrict__ w,
    const float* __restrict__ bias,
    const float* __restrict__ scale,
    const float* __restrict__ shift,
    float* __restrict__ out,
    int total_tiles)
{
    __shared__ float s_x[CI * SXD];
    __shared__ float s_w[CI * CO * 27];
    __shared__ float s_red[NTHREADS][PITCH];
    __shared__ float s_sc[CO];
    __shared__ float s_sh[CO];

    const int tid = threadIdx.x;
    for (int i = tid; i < CI * CO * 27; i += NTHREADS) s_w[i] = w[i];
    if (tid < CO) { s_sc[tid] = scale[tid]; s_sh[tid] = shift[tid]; }

    const int cls = tid >> 4, r = tid & 15;
    const int pd = c_pd[cls], ph = c_ph[cls], pw = c_pw[cls];
    const int hd = r >> 2, hh = (r >> 1) & 1, hw = r & 1;
    const int vd = 2 * hd + pd, vh = 2 * hh + ph, vw = 2 * hw + pw;
    const int vid = vd * 16 + vh * 4 + vw;

    int ad[2], ah[2], aw[2], kd[2], kh[2], kw[2];
    const int nd = setup_taps(vd, ad, kd);
    const int nh = setup_taps(vh, ah, kh);
    const int nw = setup_taps(vw, aw, kw);
    __syncthreads();

    const int TPI = 8 * ODIM * ODIM;
    for (int t = blockIdx.x; t < total_tiles; t += gridDim.x) {
        const int n   = t / TPI;
        int rem       = t - n * TPI;
        const int td  = rem / (ODIM * ODIM);
        rem          %= ODIM * ODIM;
        const int th  = rem / ODIM;
        const int tw  = rem % ODIM;
        const bool gd = (td == 7);

        const float* xb = x + (long)n * CI * IDIM * IDIM * IDIM;
        if (!gd) {
            #pragma unroll
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                s_x[i] = xb[((ic * IDIM + 4 * td + a) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
            }
        } else {
            for (int i = tid; i < CI * SXD; i += NTHREADS) {
                const int ic = i / SXD;
                const int a  = (i % SXD) / 9;
                const int b  = (i % 9) / 3;
                const int c  = i % 3;
                const int gi = 4 * td + a;
                float v = 0.0f;
                if (gi < IDIM)
                    v = xb[((ic * IDIM + gi) * IDIM + 2 * th + b) * IDIM + 2 * tw + c];
                s_x[i] = v;
            }
        }
        __syncthreads();

        float acc[CO];
        if (!gd || vd < 4) {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = bias[oc];
            conv_acc(s_x, s_w, ad, ah, aw, kd, kh, kw, nd, nh, nw, acc);
        } else {
            #pragma unroll
            for (int oc = 0; oc < CO; ++oc) acc[oc] = 0.0f;
        }

        #pragma unroll
        for (int oc = 0; oc < CO; ++oc) s_red[vid][oc] = acc[oc];
        __syncthreads();

        if (tid < 32) {
            const int oc = tid & 15, half = tid >> 4;
            const int bd4 = 2 * td + half;
            if (bd4 < ODIM) {
                float s = 0.0f;
                #pragma unroll 8
                for (int i = 0; i < 64; ++i) s += s_red[half * 64 + i][oc];
                const long oidx = (((long)n * CO + oc) * ODIM + bd4) * (ODIM * ODIM)
                                + (long)th * ODIM + tw;
                out[oidx] = s * s_sc[oc] + s_sh[oc];
            }
        }
    }
}

// ---------------------------------------------------------------------------
// BatchNorm finalization: computes mean/var from the accumulated moments,
// updates running statistics, and produces the fused scale/shift that also
// folds the 1/64 average-pool factor.
// ---------------------------------------------------------------------------
__global__ void bn_finalize_kernel(
    const double* __restrict__ gsum,
    const double* __restrict__ gsq,
    const float* __restrict__ gamma,
    const float* __restrict__ beta,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    long long* __restrict__ nbt,
    float* __restrict__ scale,
    float* __restrict__ shift,
    double m, double momentum, double eps, int track)
{
    const int c = threadIdx.x;
    if (c >= CO) return;
    const double mean = gsum[c] / m;
    double var = gsq[c] / m - mean * mean;
    if (var < 0.0) var = 0.0;
    const double invstd = rsqrt(var + eps);
    if (track) {
        const double var_unb = var * (m / (m - 1.0));
        running_mean[c] = (float)((double)running_mean[c] * (1.0 - momentum) + momentum * mean);
        running_var[c]  = (float)((double)running_var[c] * (1.0 - momentum) + momentum * var_unb);
        if (c == 0) nbt[0] += 1;
    }
    const double g = (double)gamma[c];
    scale[c] = (float)(g * invstd * (1.0 / 64.0));
    shift[c] = (float)((double)beta[c] - mean * g * invstd);
}

// ---------------------------------------------------------------------------
// Host wrappers
// ---------------------------------------------------------------------------
torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                          torch::Tensor gsum, torch::Tensor gsq) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    const int total = N * 8 * 16 * 16;
    int blocks = total < 4096 ? total : 4096;
    fused_train_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        out.data_ptr<float>(),
        gsum.data_ptr<double>(), gsq.data_ptr<double>(), total);
    return out;
}

torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias,
                         torch::Tensor scale, torch::Tensor shift) {
    const int N = x.size(0);
    auto out = torch::empty({N, CO, ODIM, ODIM, ODIM}, x.options());
    const int total = N * 8 * ODIM * ODIM;
    int blocks = total < 4096 ? total : 4096;
    fused_eval_kernel<<<blocks, NTHREADS>>>(
        x.data_ptr<float>(), w.data_ptr<float>(), bias.data_ptr<float>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        out.data_ptr<float>(), total);
    return out;
}

void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma,
                 torch::Tensor beta, torch::Tensor running_mean,
                 torch::Tensor running_var, torch::Tensor nbt,
                 torch::Tensor scale, torch::Tensor shift,
                 double m, double momentum, double eps, int64_t track) {
    bn_finalize_kernel<<<1, CO>>>(
        gsum.data_ptr<double>(), gsq.data_ptr<double>(),
        gamma.data_ptr<float>(), beta.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        (long long*)nbt.data_ptr<int64_t>(),
        scale.data_ptr<float>(), shift.data_ptr<float>(),
        m, momentum, eps, (int)track);
}
"""

cpp_source = (
    "torch::Tensor fused_train(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor gsum, torch::Tensor gsq);\n"
    "torch::Tensor fused_eval(torch::Tensor x, torch::Tensor w, torch::Tensor bias, "
    "torch::Tensor scale, torch::Tensor shift);\n"
    "void bn_finalize(torch::Tensor gsum, torch::Tensor gsq, torch::Tensor gamma, "
    "torch::Tensor beta, torch::Tensor running_mean, torch::Tensor running_var, "
    "torch::Tensor nbt, torch::Tensor scale, torch::Tensor shift, "
    "double m, double momentum, double eps, int64_t track);"
)

_ext = load_inline(
    name="fused_convT_bn_pool_v2",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["fused_train", "fused_eval", "bn_finalize"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model: ConvTranspose3d + BatchNorm3d + two AvgPool3d
    layers fused into custom CUDA kernels.  A single kernel handles the whole
    forward pass in inference mode; training mode uses one fused kernel for
    conv+pool+statistics, a tiny finalization kernel for the BatchNorm
    moments, and one elementwise scale.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias_shape):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
                                                 stride=stride, padding=padding)
        self.batch_norm = nn.BatchNorm3d(out_channels)
        self.avg_pool1 = nn.AvgPool3d(kernel_size=2)
        self.avg_pool2 = nn.AvgPool3d(kernel_size=2)
        self.ext = _ext

    def _fallback(self, x):
        x = self.conv_transpose(x)
        x = self.batch_norm(x)
        x = self.avg_pool1(x)
        x = self.avg_pool2(x)
        return x

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

        supported = (
            x.is_cuda and x.dtype == torch.float32 and
            ct.in_channels == 3 and ct.out_channels == 16 and
            ct.kernel_size == (3, 3, 3) and ct.stride == (2, 2, 2) and
            ct.padding == (1, 1, 1) and ct.output_padding == (0, 0, 0) and
            ct.dilation == (1, 1, 1) and ct.groups == 1 and ct.bias is not None and
            tuple(x.shape[-3:]) == (32, 32, 32) and
            bn.affine and bn.momentum is not None and
            (self.training or bn.track_running_stats)
        )
        if not supported:
            return self._fallback(x)

        x = x.contiguous()
        w = ct.weight.contiguous()
        b = ct.bias.contiguous()
        gamma = bn.weight.contiguous()
        beta = bn.bias.contiguous()
        dev = x.device

        if self.training:
            n = x.shape[0]
            m_count = float(n) * 250047.0  # n * 63^3, exactly representable
            gsum = torch.zeros(16, dtype=torch.float64, device=dev)
            gsq = torch.zeros(16, dtype=torch.float64, device=dev)
            raw = self.ext.fused_train(x, w, b, gsum, gsq)

            scale = torch.empty(16, dtype=torch.float32, device=dev)
            shift = torch.empty(16, dtype=torch.float32, device=dev)
            track = bn.track_running_stats
            rm = bn.running_mean if track else gamma
            rv = bn.running_var if track else gamma
            nbt = bn.num_batches_tracked if track else bn.num_batches_tracked
            self.ext.bn_finalize(gsum, gsq, gamma, beta, rm, rv, nbt,
                                 scale, shift, m_count, float(bn.momentum),
                                 float(bn.eps), 1 if track else 0)

            return torch.addcmul(shift.view(1, 16, 1, 1, 1), raw,
                                 scale.view(1, 16, 1, 1, 1))
        else:
            scale = gamma * torch.rsqrt(bn.running_var + bn.eps) * (1.0 / 64.0)
            shift = beta - bn.running_mean * (gamma * torch.rsqrt(bn.running_var + bn.eps))
            return self.ext.fused_eval(x, w, b, scale.contiguous(), shift.contiguous())


batch_size = 64
in_channels = 3
out_channels = 16
depth, height, width = 32, 32, 32
kernel_size = 3
stride = 2
padding = 1
bias_shape = (out_channels, 1, 1, 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, bias_shape]
```

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.
