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

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


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

        ``` 

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

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


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

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


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


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

        ``` 

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

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

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

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

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

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

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

    return out;
}
"""

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

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


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

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

        ``` 

        
    You are given the following architecture: 

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

import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Model that performs a 3D transposed convolution, followed by two max pooling layers and a sum operation.
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(Model, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)

    def forward(self, x):
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True) 
        return x

batch_size = 16
in_channels = 32
out_channels = 64
depth, height, width = 32, 32, 32
kernel_size = 5
stride = 2
padding = 2

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_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <float.h>

#define TILE 8
#define OC_TILE 8
#define IC_CHUNK 4
#define K 5
#define NTAP 125

// ---------------------------------------------------------------------------
// Custom ConvTranspose3d (gather formulation) for stride=2, kernel=5.
// out[b, oc, D, H, W] = sum_{ic,kd,kh,kw} X[b,ic,(D+pad-kd)/2,...] * W[ic,oc,kd,kh,kw]
// Only the parity-matched kernel taps are evaluated (saves ~7x FLOPs vs dense).
// Weight must be pre-packed to layout [IC][125][OC] (tap-major, oc contiguous).
// ---------------------------------------------------------------------------
__global__ void convt3d_s2k5_kernel(
    const float* __restrict__ X,
    const float* __restrict__ Wp,
    const float* __restrict__ bias,
    float* __restrict__ Y,
    int IC, int Di, int Hi, int Wi,
    int OC, int Do, int Ho, int Wo, int pad)
{
    __shared__ __align__(16) float sx[IC_CHUNK * 216];
    __shared__ __align__(16) float sw[NTAP * IC_CHUNK * OC_TILE];

    const int tid = threadIdx.x;                 // 0..511
    const int tx = tid & 7;
    const int ty = (tid >> 3) & 7;
    const int tz = tid >> 6;

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int tw = blockIdx.x % ntw;
    const int th = (blockIdx.x / ntw) % nth;
    const int td = blockIdx.x / (ntw * nth);

    const int W0 = tw * TILE, H0 = th * TILE, D0 = td * TILE;
    const int oc0 = blockIdx.y * OC_TILE;
    const int b   = blockIdx.z;

    const int W = W0 + tx, H = H0 + ty, D = D0 + tz;
    const bool valid_out = (D < Do) && (H < Ho) && (W < Wo);

    // input region required by this output tile
    const int d_lo = max(0, (D0 + pad - (K - 1)) / 2);
    const int d_hi = min(Di - 1, (D0 + TILE - 1 + pad) / 2);
    const int h_lo = max(0, (H0 + pad - (K - 1)) / 2);
    const int h_hi = min(Hi - 1, (H0 + TILE - 1 + pad) / 2);
    const int w_lo = max(0, (W0 + pad - (K - 1)) / 2);
    const int w_hi = min(Wi - 1, (W0 + TILE - 1 + pad) / 2);
    const int region_d = d_hi - d_lo + 1;
    const int region_h = h_hi - h_lo + 1;
    const int region_w = w_hi - w_lo + 1;

    // per-thread local tap indices (-1 = invalid)
    int dloc[K], hloc[K], wloc[K];
    #pragma unroll
    for (int kd = 0; kd < K; ++kd) {
        int t = D + pad - kd;
        int d = t >> 1;
        dloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= d_lo && d <= d_hi) ? (d - d_lo) : -1;
        t = H + pad - kd;
        d = t >> 1;
        hloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= h_lo && d <= h_hi) ? (d - h_lo) : -1;
        t = W + pad - kd;
        d = t >> 1;
        wloc[kd] = (valid_out && t >= 0 && !(t & 1) && d >= w_lo && d <= w_hi) ? (d - w_lo) : -1;
    }

    float acc[OC_TILE];
    #pragma unroll
    for (int o = 0; o < OC_TILE; ++o) acc[o] = 0.f;

    for (int ic0 = 0; ic0 < IC; ic0 += IC_CHUNK) {
        // stage input tile into shared memory
        #pragma unroll
        for (int idx = tid; idx < IC_CHUNK * 216; idx += 512) {
            int ic = idx / 216;
            int r  = idx % 216;
            int dz = r / 36;
            int dy = (r / 6) % 6;
            int dx = r % 6;
            float v = 0.f;
            if (dz < region_d && dy < region_h && dx < region_w) {
                const long xb = ((long)(b * IC + ic0 + ic) * Di + (d_lo + dz)) * Hi * Wi;
                v = X[xb + (h_lo + dy) * Wi + (w_lo + dx)];
            }
            sx[idx] = v;
        }
        // stage packed weights into shared memory
        #pragma unroll
        for (int idx = tid; idx < NTAP * IC_CHUNK * OC_TILE; idx += 512) {
            int oc  = idx & (OC_TILE - 1);
            int ic  = (idx / OC_TILE) & (IC_CHUNK - 1);
            int tap = idx / (OC_TILE * IC_CHUNK);
            sw[idx] = Wp[((long)(ic0 + ic) * NTAP + tap) * OC + (oc0 + oc)];
        }
        __syncthreads();

        if (valid_out) {
            #pragma unroll
            for (int kd = 0; kd < K; ++kd) {
                const int dl = dloc[kd];
                if (dl < 0) continue;
                #pragma unroll
                for (int kh = 0; kh < K; ++kh) {
                    const int hl = hloc[kh];
                    if (hl < 0) continue;
                    #pragma unroll
                    for (int kw = 0; kw < K; ++kw) {
                        const int wl = wloc[kw];
                        if (wl < 0) continue;
                        const int base  = (dl * 6 + hl) * 6 + wl;
                        const int tapid = kd * 25 + kh * 5 + kw;
                        #pragma unroll
                        for (int ic = 0; ic < IC_CHUNK; ++ic) {
                            const float xv = sx[ic * 216 + base];
                            const float4* wp4 = reinterpret_cast<const float4*>(&sw[(tapid * IC_CHUNK + ic) * OC_TILE]);
                            const float4 w0 = wp4[0];
                            const float4 w1 = wp4[1];
                            acc[0] = fmaf(xv, w0.x, acc[0]);
                            acc[1] = fmaf(xv, w0.y, acc[1]);
                            acc[2] = fmaf(xv, w0.z, acc[2]);
                            acc[3] = fmaf(xv, w0.w, acc[3]);
                            acc[4] = fmaf(xv, w1.x, acc[4]);
                            acc[5] = fmaf(xv, w1.y, acc[5]);
                            acc[6] = fmaf(xv, w1.z, acc[6]);
                            acc[7] = fmaf(xv, w1.w, acc[7]);
                        }
                    }
                }
            }
        }
        __syncthreads();
    }

    if (valid_out) {
        const long yb = ((long)b * OC + oc0) * Do * Ho * Wo + (long)D * Ho * Wo + (long)H * Wo + W;
        const long ostride = (long)Do * Ho * Wo;
        #pragma unroll
        for (int o = 0; o < OC_TILE; ++o) {
            Y[yb + (long)o * ostride] = acc[o] + bias[oc0 + o];
        }
    }
}

// ---------------------------------------------------------------------------
// Fused MaxPool(2)+MaxPool(3)+sum over channels.
// Composition of non-overlapping 2x2x2 then 3x3x3 max-pools equals a single
// non-overlapping 6x6x6 max-pool. Kernel computes, per output element, the
// per-channel max over the 6x6x6 window and sums over all 64 channels.
// One block (256 threads) per output element.
// ---------------------------------------------------------------------------
__global__ void pool6_sum_kernel(
    const float* __restrict__ Y,
    float* __restrict__ O,
    int C, int D, int H, int W,
    int OD, int OH, int OW)
{
    const int oid = blockIdx.x;
    const int ow = oid % OW;
    const int oh = (oid / OW) % OH;
    const int od = oid / (OW * OH);
    const int b  = blockIdx.y;

    const int tid  = threadIdx.x;   // 0..255
    const int c    = tid & 63;      // channel
    const int part = tid >> 6;      // 0..3

    const long HW = (long)H * W;
    const float* yb = Y + ((long)b * C + c) * D * HW + (long)6 * od * HW + (long)6 * oh * W + 6 * ow;

    float m = -FLT_MAX;
    for (int v = part; v < 216; v += 4) {
        const int dz = v / 36;
        const int dy = (v / 6) % 6;
        const int dx = v % 6;
        m = fmaxf(m, yb[(long)dz * HW + (long)dy * W + dx]);
    }

    __shared__ float smax[4][64];
    __shared__ float warpsum[2];
    smax[part][c] = m;
    __syncthreads();

    if (tid < 64) {
        float mm = fmaxf(fmaxf(smax[0][tid], smax[1][tid]),
                         fmaxf(smax[2][tid], smax[3][tid]));
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            mm += __shfl_down_sync(0xffffffffu, mm, off);
        if ((tid & 31) == 0) warpsum[tid >> 5] = mm;
    }
    __syncthreads();
    if (tid == 0) {
        O[((long)b * OD + od) * OH * OW + (long)oh * OW + ow] = warpsum[0] + warpsum[1];
    }
}

torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias,
                                 int64_t stride, int64_t pad)
{
    const int B  = x.size(0);
    const int IC = x.size(1);
    const int Di = x.size(2);
    const int Hi = x.size(3);
    const int Wi = x.size(4);
    const int OC = wp.size(2);

    const int Do = (Di - 1) * (int)stride - 2 * (int)pad + K;
    const int Ho = (Hi - 1) * (int)stride - 2 * (int)pad + K;
    const int Wo = (Wi - 1) * (int)stride - 2 * (int)pad + K;

    auto Y = torch::empty({B, OC, Do, Ho, Wo}, x.options());

    const int ntw = (Wo + TILE - 1) / TILE;
    const int nth = (Ho + TILE - 1) / TILE;
    const int ntd = (Do + TILE - 1) / TILE;

    dim3 grid(ntw * nth * ntd, OC / OC_TILE, B);
    convt3d_s2k5_kernel<<<grid, 512>>>(
        x.data_ptr<float>(), wp.data_ptr<float>(), bias.data_ptr<float>(),
        Y.data_ptr<float>(), IC, Di, Hi, Wi, OC, Do, Ho, Wo, (int)pad);

    // fused pool1 (k=2,s=2) + pool2 (k=3,s=3) + channel sum
    const int D1 = (Do - 2) / 2 + 1;
    const int H1 = (Ho - 2) / 2 + 1;
    const int W1 = (Wo - 2) / 2 + 1;
    const int OD = (D1 - 3) / 3 + 1;
    const int OH = (H1 - 3) / 3 + 1;
    const int OW = (W1 - 3) / 3 + 1;

    auto O = torch::empty({B, 1, OD, OH, OW}, x.options());
    dim3 pgrid(OD * OH * OW, B, 1);
    pool6_sum_kernel<<<pgrid, 256>>>(
        Y.data_ptr<float>(), O.data_ptr<float>(), OC, Do, Ho, Wo, OD, OH, OW);

    return O;
}
"""

_cpp_src = "torch::Tensor fused_forward_cuda(torch::Tensor x, torch::Tensor wp, torch::Tensor bias, int64_t stride, int64_t pad);"

_fused = load_inline(
    name="convt3d_pool_sum_fused",
    cpp_sources=_cpp_src,
    cuda_sources=_cuda_src,
    functions=["fused_forward_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized Model:
      - ConvTranspose3d replaced by a custom tiled CUDA kernel that evaluates only
        parity-matched kernel taps (stride=2, k=5) with shared-memory blocking.
      - MaxPool3d(2) + MaxPool3d(3) + channel-sum fused into a single CUDA kernel
        exploiting that the composed pooling equals one non-overlapping 6x6x6 max-pool.
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super(ModelNew, self).__init__()
        self.conv_transpose = nn.ConvTranspose3d(
            in_channels, out_channels, kernel_size, stride=stride, padding=padding
        )
        self.max_pool1 = nn.MaxPool3d(kernel_size=2)
        self.max_pool2 = nn.MaxPool3d(kernel_size=3)
        self._packed_w = None
        self._packed_version = -1

    def _get_packed_weight(self, w):
        ver = w._version
        if (
            self._packed_w is None
            or self._packed_version != ver
            or self._packed_w.device != w.device
        ):
            # [IC, OC, K, K, K] -> [IC, K*K*K, OC]
            self._packed_w = w.permute(0, 2, 3, 4, 1).contiguous()
            self._packed_version = ver
        return self._packed_w

    def forward(self, x):
        ct = self.conv_transpose
        if (
            x.is_cuda
            and x.dtype == torch.float32
            and x.is_contiguous()
            and ct.kernel_size == (5, 5, 5)
            and ct.stride == (2, 2, 2)
            and ct.padding == (2, 2, 2)
            and ct.dilation == (1, 1, 1)
            and ct.output_padding == (0, 0, 0)
            and ct.groups == 1
            and ct.bias is not None
            and ct.in_channels % 4 == 0
            and ct.out_channels % 8 == 0
            and ct.out_channels == 64
            and ct.weight.is_cuda
            and ct.weight.dtype == torch.float32
        ):
            wp = self._get_packed_weight(ct.weight)
            return _fused.fused_forward_cuda(
                x, wp, ct.bias.contiguous(), ct.stride[0], ct.padding[0]
            )
        # generic fallback
        x = self.conv_transpose(x)
        x = self.max_pool1(x)
        x = self.max_pool2(x)
        x = torch.sum(x, dim=1, keepdim=True)
        return x
```
Result of evaluating it:
The solution compiled but failed at run time. Error output:
```
CUDA error: invalid argument
Search for `cudaErrorInvalidValue' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.

```

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.
