```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 fused_transpose_center_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    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>
#include <mma.h>

namespace {

constexpr int CHANNELS = 32;
constexpr int GROUPS = 8;
constexpr int ROWS = 64;

__device__ __forceinline__ float4 zero4() {
    return make_float4(0.f, 0.f, 0.f, 0.f);
}

__device__ __forceinline__ float4 add4(float4 a, float4 b) {
    return make_float4(a.x+b.x, a.y+b.y, a.z+b.z, a.w+b.w);
}

__device__ __forceinline__ float4 mul4(float4 a, float s) {
    return make_float4(a.x*s, a.y*s, a.z*s, a.w*s);
}

__device__ __forceinline__ float4 shuffle4(float4 v, int delta) {
    return make_float4(
        __shfl_down_sync(0xffffffffu, v.x, delta),
        __shfl_down_sync(0xffffffffu, v.y, delta),
        __shfl_down_sync(0xffffffffu, v.z, delta),
        __shfl_down_sync(0xffffffffu, v.w, delta));
}

__global__ void transpose16(
    const float* __restrict__ input,
    float4* __restrict__ output,
    int spatial)
{
    __shared__ float tile[16][65];

    const int t = threadIdx.x;
    const int channel = t >> 4;
    const int position = (t & 15) * 4;
    const int begin = blockIdx.x * 64;
    const int64_t base = int64_t(blockIdx.y) * spatial * 16;

    float4 v = zero4();
    if (begin + position < spatial) {
        v = *reinterpret_cast<const float4*>(
            input + base + channel * spatial + begin + position);
    }

    tile[channel][position] = v.x;
    tile[channel][position+1] = v.y;
    tile[channel][position+2] = v.z;
    tile[channel][position+3] = v.w;
    __syncthreads();

    const int s = t >> 2;
    const int c = (t & 3) * 4;
    if (begin + s < spatial) {
        output[base / 4 + int64_t(begin+s) * 4 + (t & 3)] =
            make_float4(
                tile[c][s], tile[c+1][s],
                tile[c+2][s], tile[c+3][s]);
    }
}

// A stride-two, kernel-three transpose convolution has eight independent
// spatial parity classes. Each class is a dense TF32 matrix product.
template<int Phase, bool Training>
__global__ void convolution_stats(
    const float* __restrict__ input,
    const float* __restrict__ weight,
    float* __restrict__ output,
    float* __restrict__ sums,
    float* __restrict__ squares,
    int depth,
    int height,
    int width,
    int phase_positions,
    int parts,
    int part_offset,
    int batch)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;

    constexpr int PD = (Phase >> 2) & 1;
    constexpr int PH = (Phase >> 1) & 1;
    constexpr int PW = Phase & 1;
    constexpr int K = 16 << (PD + PH + PW);
    constexpr int AB_SIZE = (ROWS + CHANNELS) * K;
    constexpr int SHARED_SIZE =
        AB_SIZE > ROWS * CHANNELS ? AB_SIZE : ROWS * CHANNELS;

    __shared__ __align__(32) float shared[SHARED_SIZE];

    float* matrix_a = shared;
    float* matrix_b = shared + ROWS * K;

    const int t = threadIdx.x;
    const int warp = t >> 5;
    const int n = blockIdx.y;
    const int start = blockIdx.x * ROWS;
    const int phase_width = width - PW;
    const int phase_height = height - PH;
    const int input_spatial = depth * height * width;
    const int output_height = height * 2 - 1;
    const int output_width = width * 2 - 1;
    const int output_spatial =
        (depth * 2 - 1) * output_height * output_width;

    for (int index = t; index < ROWS * K; index += 256) {
        const int row = index / K;
        const int k = index % K;
        const int position = start + row;

        float value = 0.f;
        if (position < phase_positions) {
            int q = position;
            const int w = q % phase_width;
            q /= phase_width;
            const int h = q % phase_height;
            const int d = q / phase_height;

            int tap = k >> 4;
            const int kw = PW ? 2 * (tap & 1) : 1;
            if (PW) tap >>= 1;
            const int kh = PH ? 2 * (tap & 1) : 1;
            if (PH) tap >>= 1;
            const int kd = PD ? 2 * (tap & 1) : 1;

            const int iw = w + (PW && kw == 0);
            const int ih = h + (PH && kh == 0);
            const int id = d + (PD && kd == 0);

            const int64_t address =
                (int64_t(n) * input_spatial
                 + (id * height + ih) * width + iw) * 16
                + (k & 15);
            value = input[address];
        }
        matrix_a[index] = __float_to_tf32(value);
    }

    // Column-major B: one contiguous K-vector per output channel.
    for (int index = t; index < CHANNELS * K; index += 256) {
        const int c = index / K;
        const int k = index % K;
        int tap = k >> 4;

        const int kw = PW ? 2 * (tap & 1) : 1;
        if (PW) tap >>= 1;
        const int kh = PH ? 2 * (tap & 1) : 1;
        if (PH) tap >>= 1;
        const int kd = PD ? 2 * (tap & 1) : 1;

        const int address =
            (((k & 15) * 27 + kd * 9 + kh * 3 + kw) * 32) + c;
        matrix_b[index] = __float_to_tf32(weight[address]);
    }
    __syncthreads();

    wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major> a;
    wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::col_major> b;
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> accumulator;

    wmma::fill_fragment(accumulator, 0.f);

    const int row_start = (warp >> 1) * 16;
    const int channel_start = (warp & 1) * 16;

    #pragma unroll
    for (int k = 0; k < K; k += 8) {
        wmma::load_matrix_sync(a, matrix_a + row_start * K + k, K);
        wmma::load_matrix_sync(b, matrix_b + channel_start * K + k, K);
        wmma::mma_sync(accumulator, a, b, accumulator);
    }

    // Reuse the operand storage for the output tile.
    __syncthreads();
    wmma::store_matrix_sync(
        shared + row_start * CHANNELS + channel_start,
        accumulator, CHANNELS, wmma::mem_row_major);
    __syncthreads();

    const int c = t & 31;
    const int first_row = t >> 5;
    float sum = 0.f;
    float square = 0.f;

    #pragma unroll
    for (int row = first_row; row < ROWS; row += 8) {
        const int position = start + row;
        if (position < phase_positions) {
            int q = position;
            const int w = q % phase_width;
            q /= phase_width;
            const int h = q % phase_height;
            const int d = q / phase_height;

            const float value = shared[row * CHANNELS + c];
            const int output_position =
                ((2*d + PD) * output_height + 2*h + PH)
                * output_width + 2*w + PW;

            output[
                (int64_t(n) * output_spatial + output_position) * 32 + c
            ] = value;

            sum += value;
            if (Training)
                square = fmaf(value, value, square);
        }
    }

    __syncthreads();
    shared[t] = sum;
    if (Training)
        shared[256+t] = square;
    __syncthreads();

    if (t < 32) {
        sum = shared[t];
        if (Training)
            square = shared[256+t];

        #pragma unroll
        for (int w = 1; w < 8; ++w) {
            sum += shared[w*32+t];
            if (Training)
                square += shared[256+w*32+t];
        }

        const int64_t dst =
            ((int64_t(t >> 2) * batch + n) * parts
             + part_offset + blockIdx.x) * 4 + (t & 3);
        sums[dst] = sum;
        if (Training)
            squares[dst] = square;
    }
#endif
}

template<bool Training>
__global__ void reduce_sample_stats(
    const float4* __restrict__ sums,
    const float4* __restrict__ squares,
    float4* __restrict__ means,
    float4* __restrict__ seconds,
    int parts,
    int spatial,
    int batch)
{
    const int group = blockIdx.x;
    const int n = blockIdx.y;
    const int t = threadIdx.x;
    const int64_t base = (int64_t(group) * batch + n) * parts;

    float4 sum = zero4();
    float4 square = zero4();

    for (int p = t; p < parts; p += 256) {
        sum = add4(sum, sums[base+p]);
        if (Training)
            square = add4(square, squares[base+p]);
    }

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum = add4(sum, shuffle4(sum, offset));
        if (Training)
            square = add4(square, shuffle4(square, offset));
    }

    __shared__ float4 warp_sums[8];
    __shared__ float4 warp_squares[8];

    if ((t & 31) == 0) {
        warp_sums[t >> 5] = sum;
        if (Training)
            warp_squares[t >> 5] = square;
    }
    __syncthreads();

    if (t < 32) {
        sum = t < 8 ? warp_sums[t] : zero4();
        if (Training)
            square = t < 8 ? warp_squares[t] : zero4();

        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum = add4(sum, shuffle4(sum, offset));
            if (Training)
                square = add4(square, shuffle4(square, offset));
        }

        if (t == 0) {
            const float inv = 1.f / float(spatial);
            means[int64_t(n) * GROUPS + group] = mul4(sum, inv);
            if (Training)
                seconds[int64_t(n) * GROUPS + group] =
                    mul4(square, inv);
        }
    }
}

template<bool Training>
__global__ void finish_stats(
    const float* __restrict__ means,
    const float* __restrict__ seconds,
    const float* __restrict__ gamma,
    const float* __restrict__ conv_bias,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    float* __restrict__ scales,
    int batch,
    int spatial,
    float eps,
    float momentum)
{
    const int c = threadIdx.x;
    float variance;

    if (Training) {
        float mean = 0.f;
        float second = 0.f;

        for (int n = 0; n < batch; ++n) {
            mean += means[n*32+c];
            second += seconds[n*32+c];
        }

        mean /= float(batch);
        second /= float(batch);
        variance = fmaxf(fmaf(-mean, mean, second), 0.f);

        const float shifted_mean =
            mean + (conv_bias ? conv_bias[c] : 0.f);
        const float count = float(batch) * float(spatial);

        running_mean[c] =
            (1.f-momentum) * running_mean[c] + momentum * shifted_mean;
        running_var[c] =
            (1.f-momentum) * running_var[c]
            + momentum * variance * (count / (count-1.f));
    } else {
        variance = running_var[c];
    }

    scales[c] = gamma[c] * rsqrtf(variance + eps);
}

__device__ __forceinline__ float4 load_stream(const float4* p) {
    float4 value;
    asm volatile(
        "ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(value.x), "=f"(value.y), "=f"(value.z), "=f"(value.w)
        : "l"(p));
    return value;
}

__device__ __forceinline__ void store_stream(float4* p, float4 v) {
    asm volatile(
        "st.global.cs.v4.f32 [%0], {%1, %2, %3, %4};"
        :
        : "l"(p), "f"(v.x), "f"(v.y), "f"(v.z), "f"(v.w)
        : "memory");
}

__global__ void center_output(
    float4* output,
    const float4* __restrict__ means,
    const float4* __restrict__ scales,
    int64_t sample_vectors)
{
    const int n = blockIdx.y;
    const int group = threadIdx.x & 7;
    const float4 mean = means[int64_t(n) * GROUPS + group];
    const float4 scale = scales[group];

    const int64_t begin = int64_t(blockIdx.x) * 1024 + threadIdx.x;
    const int64_t base = int64_t(n) * sample_vectors;

    #pragma unroll
    for (int k = 0; k < 1024; k += 256) {
        const int64_t index = begin + k;
        if (index < sample_vectors) {
            const float4 v = load_stream(output + base + index);
            store_stream(
                output + base + index,
                make_float4(
                    (v.x-mean.x)*scale.x,
                    (v.y-mean.y)*scale.y,
                    (v.z-mean.z)*scale.z,
                    (v.w-mean.w)*scale.w));
        }
    }
}

template<int Phase, bool Training>
void launch_phase(
    const float* input,
    const float* weight,
    float* output,
    float* sums,
    float* squares,
    int depth,
    int height,
    int width,
    int parts,
    int batch,
    int& offset,
    cudaStream_t stream)
{
    constexpr int PD = (Phase >> 2) & 1;
    constexpr int PH = (Phase >> 1) & 1;
    constexpr int PW = Phase & 1;

    const int positions =
        (depth-PD) * (height-PH) * (width-PW);
    const int blocks = (positions + ROWS - 1) / ROWS;

    convolution_stats<Phase, Training>
        <<<dim3(blocks, batch), 256, 0, stream>>>(
            input, weight, output, sums, squares,
            depth, height, width, positions,
            parts, offset, batch);

    offset += blocks;
}

template<bool Training>
void launch_all(
    const float* input,
    const float* weight,
    float* output,
    float* sums,
    float* squares,
    float* means,
    float* seconds,
    float* scales,
    const float* gamma,
    const float* conv_bias,
    float* running_mean,
    float* running_var,
    int depth,
    int height,
    int width,
    int parts,
    int batch,
    int spatial,
    float eps,
    float momentum,
    cudaStream_t stream)
{
    int offset = 0;

    #define LAUNCH_PHASE(P) \
        launch_phase<P, Training>( \
            input, weight, output, sums, squares, \
            depth, height, width, parts, batch, offset, stream)

    LAUNCH_PHASE(0);
    LAUNCH_PHASE(1);
    LAUNCH_PHASE(2);
    LAUNCH_PHASE(3);
    LAUNCH_PHASE(4);
    LAUNCH_PHASE(5);
    LAUNCH_PHASE(6);
    LAUNCH_PHASE(7);

    #undef LAUNCH_PHASE

    reduce_sample_stats<Training>
        <<<dim3(GROUPS, batch), 256, 0, stream>>>(
            reinterpret_cast<const float4*>(sums),
            reinterpret_cast<const float4*>(squares),
            reinterpret_cast<float4*>(means),
            reinterpret_cast<float4*>(seconds),
            parts, spatial, batch);

    finish_stats<Training><<<1, 32, 0, stream>>>(
        means, seconds, gamma, conv_bias,
        running_mean, running_var, scales,
        batch, spatial, eps, momentum);

    const int64_t vectors = int64_t(spatial) * GROUPS;
    center_output
        <<<dim3((vectors+1023)/1024, batch), 256, 0, stream>>>(
            reinterpret_cast<float4*>(output),
            reinterpret_cast<const float4*>(means),
            reinterpret_cast<const float4*>(scales),
            vectors);
}

} // namespace

torch::Tensor fused_transpose_center_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor conv_bias,
    torch::Tensor gamma,
    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.size(1) == 16,
                "Expected NCDHW input with 16 channels");
    TORCH_CHECK(weight.is_contiguous(at::MemoryFormat::ChannelsLast3d),
                "Expected channels-last convolution weights");

    const c10::cuda::CUDAGuard guard(x.device());
    const auto stream = at::cuda::getCurrentCUDAStream();

    const int batch = x.size(0);
    const int depth = x.size(2);
    const int height = x.size(3);
    const int width = x.size(4);
    const int input_spatial = depth * height * width;
    const int od = depth * 2 - 1;
    const int oh = height * 2 - 1;
    const int ow = width * 2 - 1;
    const int spatial = od * oh * ow;

    if (!x.is_contiguous(at::MemoryFormat::ChannelsLast3d)) {
        if (x.is_contiguous() && input_spatial % 4 == 0
            && x.storage_offset() % 4 == 0) {
            auto converted = torch::empty(
                x.sizes(),
                x.options().memory_format(at::MemoryFormat::ChannelsLast3d));

            transpose16
                <<<dim3((input_spatial+63)/64, batch), 256, 0, stream>>>(
                    x.data_ptr<float>(),
                    reinterpret_cast<float4*>(converted.data_ptr<float>()),
                    input_spatial);
            x = converted;
        } else {
            x = x.contiguous(at::MemoryFormat::ChannelsLast3d);
        }
    }

    int parts = 0;
    for (int phase = 0; phase < 8; ++phase) {
        const int positions =
            (depth-((phase >> 2) & 1))
            * (height-((phase >> 1) & 1))
            * (width-(phase & 1));
        parts += (positions + ROWS - 1) / ROWS;
    }

    auto output = torch::empty(
        {batch, 32, od, oh, ow},
        x.options().memory_format(at::MemoryFormat::ChannelsLast3d));

    const int64_t partial_size = int64_t(batch) * parts * 32;
    auto partial = torch::empty(
        {partial_size * (training ? 2 : 1)}, x.options());

    auto statistics = torch::empty(
        {int64_t(batch) * 32 * (training ? 2 : 1) + 32},
        x.options());

    float* sums = partial.data_ptr<float>();
    float* squares = training ? sums + partial_size : nullptr;
    float* means = statistics.data_ptr<float>();
    float* seconds = training ? means + int64_t(batch)*32 : nullptr;
    float* scales =
        means + int64_t(batch)*32*(training ? 2 : 1);
    const float* bias =
        conv_bias.numel() ? conv_bias.data_ptr<float>() : nullptr;

    #define RUN(TRAINING) \
        launch_all<TRAINING>( \
            x.data_ptr<float>(), weight.data_ptr<float>(), \
            output.data_ptr<float>(), sums, squares, \
            means, seconds, scales, gamma.data_ptr<float>(), bias, \
            running_mean.data_ptr<float>(), running_var.data_ptr<float>(), \
            depth, height, width, parts, batch, spatial, \
            float(eps), float(momentum), stream)

    if (training) {
        RUN(true);
    } else {
        RUN(false);
    }

    #undef RUN

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_ops = load_inline(
    name="polyphase_transpose3d_fused_stats_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["fused_transpose_center_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _FusedTransposeCenter(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx, x, weight, conv_bias, gamma, beta,
        running_mean, running_var, eps, momentum, training
    ):
        ctx.save_for_backward(
            x, weight, gamma,
            x.new_empty(0) if training else running_var.clone(),
        )
        ctx.eps = eps
        ctx.training = training
        ctx.has_conv_bias = conv_bias.numel() != 0

        return _ops.fused_transpose_center_cuda(
            x, weight, conv_bias, gamma,
            running_mean, running_var, eps, momentum, training,
        )

    @staticmethod
    def backward(ctx, grad_output):
        x, weight, gamma, running_var = ctx.saved_tensors
        higher_order = torch.is_grad_enabled()

        with torch.enable_grad():
            xx = x.detach().requires_grad_(True)
            ww = weight.detach().requires_grad_(True)
            gg = gamma.detach().requires_grad_(True)

            convolution = F.conv_transpose3d(
                xx, ww, None, stride=2, padding=1,
            )
            centered = convolution - convolution.mean(
                dim=(2, 3, 4), keepdim=True,
            )

            if ctx.training:
                variance = convolution.var(
                    dim=(0, 2, 3, 4), unbiased=False, keepdim=True,
                )
            else:
                variance = running_var.view(1, -1, 1, 1, 1)

            output = (
                centered
                * torch.rsqrt(variance + ctx.eps)
                * gg.view(1, -1, 1, 1, 1)
            )

            grad_x, grad_weight, grad_gamma = torch.autograd.grad(
                output,
                (xx, ww, gg),
                grad_output,
                create_graph=higher_order,
            )

        return (
            grad_x,
            grad_weight,
            torch.zeros_like(gamma) if ctx.has_conv_bias else None,
            grad_gamma,
            torch.zeros_like(gamma),
            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)
        self.conv_transpose.to(memory_format=torch.channels_last_3d)
        self._device_capabilities = {}

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

        eligible = (
            x.is_cuda
            and x.dtype == torch.float32
            and x.dim() == 5
            and x.size(1) == 16
            and 0 < x.size(0) <= 65535
            and min(x.shape[2:]) >= 2
            and conv.in_channels == 16
            and conv.out_channels == 32
            and conv.kernel_size == (3, 3, 3)
            and conv.stride == (2, 2, 2)
            and conv.padding == (1, 1, 1)
            and conv.output_padding == (0, 0, 0)
            and conv.dilation == (1, 1, 1)
            and conv.groups == 1
            and conv.weight.dtype == torch.float32
            and conv.weight.device == x.device
            and conv.weight.is_contiguous(
                memory_format=torch.channels_last_3d
            )
            and bn.affine
            and bn.track_running_stats
            and bn.weight.dtype == torch.float32
            and bn.running_mean.dtype == torch.float32
            and bn.running_var.dtype == torch.float32
            and bn.weight.is_contiguous()
            and bn.running_mean.is_contiguous()
            and bn.running_var.is_contiguous()
            and bn.weight.device == x.device
            and bn.running_mean.device == x.device
            and bn.running_var.device == x.device
            and not torch.is_autocast_enabled()
            and (
                conv.bias is None
                or (
                    conv.bias.dtype == torch.float32
                    and conv.bias.device == x.device
                    and conv.bias.is_contiguous()
                )
            )
        )

        if eligible:
            spatial = (
                (2 * x.size(2) - 1)
                * (2 * x.size(3) - 1)
                * (2 * x.size(4) - 1)
            )
            eligible = spatial < 2**30

        if eligible:
            device_index = x.device.index
            capability = self._device_capabilities.get(device_index)
            if capability is None:
                capability = torch.cuda.get_device_capability(x.device)[0]
                self._device_capabilities[device_index] = capability
            eligible = capability >= 8

        if not eligible:
            y = bn(conv(x))
            return y - y.mean(dim=(2, 3, 4), keepdim=True)

        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 x.new_empty(0)

        needs_grad = torch.is_grad_enabled() and (
            x.requires_grad
            or conv.weight.requires_grad
            or conv_bias.requires_grad
            or bn.weight.requires_grad
            or bn.bias.requires_grad
        )

        if needs_grad:
            return _FusedTransposeCenter.apply(
                x, conv.weight, conv_bias, bn.weight, bn.bias,
                bn.running_mean, bn.running_var,
                bn.eps, momentum, bn.training,
            )

        return _ops.fused_transpose_center_cuda(
            x, conv.weight, conv_bias, bn.weight,
            bn.running_mean, bn.running_var,
            bn.eps, momentum, bn.training,
        )
```