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

_cpp_source = r"""
torch::Tensor centered_bn_nhwc_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor conv_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps,
    double momentum,
    bool training);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

constexpr int TILE = 1024;
constexpr int CHANNEL_TILE = 32;

__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    float2* __restrict__ partial,
    int channels, int spatial, int parts)
{
    const int lane = threadIdx.x & 31;
    const int row = threadIdx.x >> 5;
    const int c = blockIdx.z * CHANNEL_TILE + lane;
    const int n = blockIdx.y;
    const int p = blockIdx.x;
    const int begin = p * TILE;
    const int end = min(begin + TILE, spatial);

    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;

    if (c < channels) {
        const float* base = x + int64_t(n) * spatial * channels + c;
        for (int i = begin + row; i < end; i += 32) {
            float a = base[int64_t(i) * channels];
            float b = i + 8 < end ? base[int64_t(i + 8) * channels] : 0.f;
            float d = i + 16 < end ? base[int64_t(i + 16) * channels] : 0.f;
            float e = i + 24 < end ? base[int64_t(i + 24) * channels] : 0.f;
            s0 += a;
            s1 += b;
            s2 += d;
            s3 += e;
            q0 = fmaf(a, a, q0);
            q1 = fmaf(b, b, q1);
            q2 = fmaf(d, d, q2);
            q3 = fmaf(e, e, q3);
        }
    }

    __shared__ float sums[8][32];
    __shared__ float squares[8][32];
    sums[row][lane] = (s0 + s1) + (s2 + s3);
    squares[row][lane] = (q0 + q1) + (q2 + q3);
    __syncthreads();

    if (row == 0 && c < channels) {
        float s = 0.f, q = 0.f;
        #pragma unroll
        for (int r = 0; r < 8; ++r) {
            s += sums[r][lane];
            q += squares[r][lane];
        }
        partial[(int64_t(n) * parts + p) * channels + c] =
            make_float2(s, q);
    }
}

__global__ void finish_stats_kernel(
    const float2* __restrict__ partial,
    const float* __restrict__ weight,
    const float* __restrict__ conv_bias,
    float* __restrict__ means,
    float* __restrict__ scales,
    float* running_mean,
    float* running_var,
    int batch, int channels, int spatial, int parts,
    float eps, float momentum, bool training)
{
    const int lane = threadIdx.x & 31;
    const int row = threadIdx.x >> 5;
    const int c = blockIdx.x * CHANNEL_TILE + lane;

    float channel_sum = 0.f;
    float channel_square = 0.f;

    if (c < channels) {
        for (int n = row; n < batch; n += 16) {
            float s0 = 0.f, s1 = 0.f;
            float q0 = 0.f, q1 = 0.f;
            int p = 0;
            for (; p + 1 < parts; p += 2) {
                float2 a = partial[
                    (int64_t(n) * parts + p) * channels + c];
                float2 b = partial[
                    (int64_t(n) * parts + p + 1) * channels + c];
                s0 += a.x;
                s1 += b.x;
                q0 += a.y;
                q1 += b.y;
            }
            if (p < parts) {
                float2 a = partial[
                    (int64_t(n) * parts + p) * channels + c];
                s0 += a.x;
                q0 += a.y;
            }

            float s = s0 + s1;
            means[n * channels + c] = s / float(spatial);
            channel_sum += s;
            channel_square += q0 + q1;
        }
    }

    __shared__ float sums[16][32];
    __shared__ float squares[16][32];
    sums[row][lane] = channel_sum;
    squares[row][lane] = channel_square;
    __syncthreads();

    if (row == 0 && c < channels) {
        float variance;
        if (training) {
            float s = 0.f, q = 0.f;
            #pragma unroll
            for (int r = 0; r < 16; ++r) {
                s += sums[r][lane];
                q += squares[r][lane];
            }
            float count = float(batch) * float(spatial);
            float mean = s / count;
            variance = fmaxf(fmaf(-mean, mean, q / count), 0.f);

            float biased_mean = mean + (conv_bias ? conv_bias[c] : 0.f);
            float unbiased = variance * (count / (count - 1.f));
            running_mean[c] =
                (1.f - momentum) * running_mean[c] + momentum * biased_mean;
            running_var[c] =
                (1.f - momentum) * running_var[c] + momentum * unbiased;
        } else {
            variance = running_var[c];
        }
        scales[c] = weight[c] * rsqrtf(variance + eps);
    }
}

__global__ void center_vector_kernel(
    const float* __restrict__ x,
    const float* __restrict__ means,
    const float* __restrict__ scales,
    float* __restrict__ out,
    int channels, int64_t sample_elements)
{
    const int n = blockIdx.y;
    const int64_t vectors = sample_elements / 4;
    const int64_t begin = int64_t(blockIdx.x) * 1024 + threadIdx.x;
    const float4* input = reinterpret_cast<const float4*>(
        x + int64_t(n) * sample_elements);
    float4* output = reinterpret_cast<float4*>(
        out + int64_t(n) * sample_elements);

    #pragma unroll
    for (int k = 0; k < 4; ++k) {
        int64_t i = begin + k * 256;
        if (i < vectors) {
            int c = (i * 4) % channels;
            float4 v = input[i];
            float4 m = *reinterpret_cast<const float4*>(
                means + n * channels + c);
            float4 s = *reinterpret_cast<const float4*>(scales + c);
            output[i] = make_float4(
                (v.x - m.x) * s.x,
                (v.y - m.y) * s.y,
                (v.z - m.z) * s.z,
                (v.w - m.w) * s.w);
        }
    }
}

__global__ void center_scalar_kernel(
    const float* __restrict__ x,
    const float* __restrict__ means,
    const float* __restrict__ scales,
    float* __restrict__ out,
    int channels, int64_t sample_elements)
{
    const int n = blockIdx.y;
    const int64_t begin = int64_t(blockIdx.x) * 4096 + threadIdx.x;
    const int64_t base = int64_t(n) * sample_elements;

    #pragma unroll
    for (int k = 0; k < 16; ++k) {
        int64_t i = begin + k * 256;
        if (i < sample_elements) {
            int c = i % channels;
            out[base + i] =
                (x[base + i] - means[n * channels + c]) * scales[c];
        }
    }
}

torch::Tensor centered_bn_nhwc_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor conv_bias,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    double eps,
    double momentum,
    bool training)
{
    TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat,
                "Expected CUDA float32 input");
    TORCH_CHECK(x.dim() == 5 &&
                x.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "Expected channels-last 3D input");

    c10::cuda::CUDAGuard guard(x.device());
    int batch = x.size(0);
    int channels = x.size(1);
    int spatial = x.size(2) * x.size(3) * x.size(4);
    TORCH_CHECK(batch > 0 && spatial > 0, "Expected nonempty input");
    TORCH_CHECK(!training || int64_t(batch) * spatial > 1,
                "Expected more than one value per channel");

    int parts = (spatial + TILE - 1) / TILE;
    auto partial = torch::empty({batch, parts, channels, 2}, x.options());
    auto means = torch::empty({batch, channels}, x.options());
    auto scales = torch::empty({channels}, x.options());
    auto out = torch::empty_like(x);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    partial_stats_kernel<<<
        dim3(parts, batch, (channels + 31) / 32), 256, 0, stream>>>(
        x.data_ptr<float>(),
        reinterpret_cast<float2*>(partial.data_ptr<float>()),
        channels, spatial, parts);

    finish_stats_kernel<<<(channels + 31) / 32, 512, 0, stream>>>(
        reinterpret_cast<const float2*>(partial.data_ptr<float>()),
        weight.data_ptr<float>(),
        conv_bias.numel() ? conv_bias.data_ptr<float>() : nullptr,
        means.data_ptr<float>(), scales.data_ptr<float>(),
        running_mean.data_ptr<float>(), running_var.data_ptr<float>(),
        batch, channels, spatial, parts,
        float(eps), float(momentum), training);

    int64_t sample_elements = int64_t(spatial) * channels;
    dim3 output_grid((sample_elements + 4095) / 4096, batch);

    if (channels % 4 == 0) {
        center_vector_kernel<<<output_grid, 256, 0, stream>>>(
            x.data_ptr<float>(), means.data_ptr<float>(),
            scales.data_ptr<float>(), out.data_ptr<float>(),
            channels, sample_elements);
    } else {
        center_scalar_kernel<<<output_grid, 256, 0, stream>>>(
            x.data_ptr<float>(), means.data_ptr<float>(),
            scales.data_ptr<float>(), out.data_ptr<float>(),
            channels, sample_elements);
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_centered_bn = load_inline(
    name="conv_transpose3d_centered_bn_nhwc_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["centered_bn_nhwc_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _CenteredBatchNorm(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, conv_bias,
                running_mean, running_var, eps, momentum, training):
        out = _centered_bn.centered_bn_nhwc_cuda(
            x, weight, conv_bias, running_mean, running_var,
            eps, momentum, training,
        )
        ctx.save_for_backward(
            x, weight, running_var if training else running_var.clone()
        )
        ctx.eps = eps
        ctx.training = training
        ctx.has_conv_bias = conv_bias.numel() != 0
        return out

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, running_var = ctx.saved_tensors
        spatial_dims = (2, 3, 4)
        channel_dims = (0, 2, 3, 4)
        shape = (1, -1, 1, 1, 1)

        centered = x - x.mean(dim=spatial_dims, keepdim=True)
        grad_centered = (
            grad_output - grad_output.mean(dim=spatial_dims, keepdim=True)
        )

        if ctx.training:
            variance, mean = torch.var_mean(
                x, dim=channel_dims, unbiased=False, keepdim=True
            )
            inv_std = torch.rsqrt(variance + ctx.eps)
        else:
            inv_std = torch.rsqrt(running_var.view(shape) + ctx.eps)

        correlation = (grad_output * centered).sum(
            dim=channel_dims, keepdim=True
        )
        grad_x = grad_centered * weight.view(shape) * inv_std

        if ctx.training:
            count = x.numel() // x.size(1)
            grad_x = grad_x - (
                (x - mean) * weight.view(shape)
                * inv_std.pow(3) * correlation / count
            )

        grad_weight = (correlation * inv_std).reshape_as(weight)
        grad_bias = torch.zeros_like(weight)
        grad_conv_bias = (
            torch.zeros_like(weight) if ctx.has_conv_bias else None
        )
        return (
            grad_x, grad_weight, grad_bias, grad_conv_bias,
            None, None, None, None, None,
        )


class ModelNew(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size,
                 stride, padding, bias=True):
        super().__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):
        conv = self.conv_transpose
        bn = self.batch_norm

        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or conv.weight.dtype != torch.float32
            or not bn.affine
            or not bn.track_running_stats
            or bn.weight.dtype != torch.float32
            or bn.running_mean.dtype != torch.float32
            or x.numel() == 0
            or x.size(0) > 65535
        ):
            y = bn(conv(x))
            return y - y.mean(dim=(2, 3, 4), keepdim=True)

        # The convolution bias and BN shift cancel during spatial centering.
        y = F.conv_transpose3d(
            x.contiguous(memory_format=torch.channels_last_3d),
            conv.weight,
            None,
            conv.stride,
            conv.padding,
            conv.output_padding,
            conv.groups,
            conv.dilation,
        )

        if y.dtype != torch.float32:
            if conv.bias is not None:
                y = y + conv.bias.to(y.dtype).view(1, -1, 1, 1, 1)
            y = bn(y)
            return y - y.mean(dim=(2, 3, 4), keepdim=True)

        y = y.contiguous(memory_format=torch.channels_last_3d)

        momentum = bn.momentum if bn.momentum is not None else 0.0
        if bn.training:
            bn.num_batches_tracked.add_(1)
            if bn.momentum is None:
                momentum = 1.0 / bn.num_batches_tracked.item()

        conv_bias = (
            conv.bias if conv.bias is not None else y.new_empty(0)
        )
        return _CenteredBatchNorm.apply(
            y, bn.weight, bn.bias, conv_bias,
            bn.running_mean, bn.running_var,
            bn.eps, momentum, bn.training,
        )
```