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


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

torch::Tensor bn_center_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches,
    double momentum,
    double eps,
    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.h>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int TILE = 4096;

__device__ __forceinline__ float warp_sum(float x) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        x += __shfl_down_sync(0xffffffff, x, offset);
    return x;
}

__device__ __forceinline__ float block_sum(float x, float* shared) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    x = warp_sum(x);
    if (lane == 0)
        shared[warp] = x;
    __syncthreads();

    x = threadIdx.x < 8 ? shared[lane] : 0.0f;
    if (warp == 0)
        x = warp_sum(x);
    return x;
}

template<bool TRAINING>
__global__ void partial_stats_kernel(
    const float* __restrict__ x,
    float* __restrict__ sums,
    float* __restrict__ squares,
    int spatial,
    int parts)
{
    const int nc = blockIdx.x / parts;
    const int part = blockIdx.x - nc * parts;
    const int begin = part * TILE;
    const int end = min(begin + TILE, spatial);
    const int64_t base = static_cast<int64_t>(nc) * spatial;

    float sum = 0.0f;
    float square = 0.0f;
    for (int i = begin + threadIdx.x; i < end; i += THREADS) {
        const float value = x[base + i];
        sum += value;
        if (TRAINING)
            square = fmaf(value, value, square);
    }

    __shared__ float sum_shared[8];
    __shared__ float square_shared[8];
    sum = block_sum(sum, sum_shared);
    if (TRAINING)
        square = block_sum(square, square_shared);

    if (threadIdx.x == 0) {
        sums[blockIdx.x] = sum;
        if (TRAINING)
            squares[blockIdx.x] = square;
    }
}

template<bool TRAINING>
__global__ void finish_stats_kernel(
    const float* __restrict__ sums,
    const float* __restrict__ squares,
    const float* __restrict__ weight,
    float* __restrict__ running_mean,
    float* __restrict__ running_var,
    int64_t* __restrict__ batches,
    float* __restrict__ instance_means,
    float* __restrict__ scales,
    int batch,
    int channels,
    int spatial,
    int parts,
    float momentum,
    float eps)
{
    const int c = blockIdx.x;

    for (int n = threadIdx.x; n < batch; n += THREADS) {
        const int offset = (n * channels + c) * parts;
        float sum = 0.0f;
        for (int p = 0; p < parts; ++p)
            sum += sums[offset + p];
        instance_means[n * channels + c] = sum / float(spatial);
    }

    if (TRAINING) {
        float sum = 0.0f;
        float square = 0.0f;
        for (int j = threadIdx.x; j < batch * parts; j += THREADS) {
            const int n = j / parts;
            const int p = j - n * parts;
            const int offset = (n * channels + c) * parts + p;
            sum += sums[offset];
            square += squares[offset];
        }

        __shared__ float sum_shared[8];
        __shared__ float square_shared[8];
        sum = block_sum(sum, sum_shared);
        square = block_sum(square, square_shared);

        if (threadIdx.x == 0) {
            const float count = float(batch) * float(spatial);
            const float mean = sum / count;
            const float variance =
                fmaxf(fmaf(-mean, mean, square / count), 0.0f);

            scales[c] = weight[c] * rsqrtf(variance + eps);
            running_mean[c] =
                (1.0f - momentum) * running_mean[c] + momentum * mean;
            running_var[c] =
                (1.0f - momentum) * running_var[c]
                + momentum * variance * (count / (count - 1.0f));

            if (c == 0)
                batches[0] += 1;
        }
    } else if (threadIdx.x == 0) {
        scales[c] = weight[c] * rsqrtf(running_var[c] + eps);
    }
}

__global__ void center_scale_kernel(
    const float* __restrict__ x,
    const float* __restrict__ means,
    const float* __restrict__ scales,
    float* __restrict__ output,
    int channels,
    int spatial,
    int parts)
{
    const int nc = blockIdx.x / parts;
    const int part = blockIdx.x - nc * parts;
    const int begin = part * TILE;
    const int end = min(begin + TILE, spatial);
    const int64_t base = static_cast<int64_t>(nc) * spatial;
    const float mean = means[nc];
    const float scale = scales[nc % channels];

    for (int i = begin + threadIdx.x; i < end; i += THREADS)
        output[base + i] = (x[base + i] - mean) * scale;
}

torch::Tensor bn_center_cuda(
    torch::Tensor x,
    torch::Tensor weight,
    torch::Tensor running_mean,
    torch::Tensor running_var,
    torch::Tensor batches,
    double momentum,
    double eps,
    bool training)
{
    const c10::cuda::CUDAGuard guard(x.device());

    const int batch = x.size(0);
    const int channels = x.size(1);
    const int spatial = x.size(2) * x.size(3) * x.size(4);
    const int parts = (spatial + TILE - 1) / TILE;
    const int blocks = batch * channels * parts;

    auto sums = torch::empty({batch, channels, parts}, x.options());
    auto squares = torch::empty_like(sums);
    auto means = torch::empty({batch, channels}, x.options());
    auto scales = torch::empty({channels}, x.options());
    auto output = torch::empty_like(x);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (training) {
        partial_stats_kernel<true><<<blocks, THREADS, 0, stream>>>(
            x.data_ptr<float>(), sums.data_ptr<float>(),
            squares.data_ptr<float>(), spatial, parts);

        finish_stats_kernel<true><<<channels, THREADS, 0, stream>>>(
            sums.data_ptr<float>(), squares.data_ptr<float>(),
            weight.data_ptr<float>(), running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(), batches.data_ptr<int64_t>(),
            means.data_ptr<float>(), scales.data_ptr<float>(),
            batch, channels, spatial, parts, float(momentum), float(eps));
    } else {
        partial_stats_kernel<false><<<blocks, THREADS, 0, stream>>>(
            x.data_ptr<float>(), sums.data_ptr<float>(),
            squares.data_ptr<float>(), spatial, parts);

        finish_stats_kernel<false><<<channels, THREADS, 0, stream>>>(
            sums.data_ptr<float>(), squares.data_ptr<float>(),
            weight.data_ptr<float>(), running_mean.data_ptr<float>(),
            running_var.data_ptr<float>(), batches.data_ptr<int64_t>(),
            means.data_ptr<float>(), scales.data_ptr<float>(),
            batch, channels, spatial, parts, float(momentum), float(eps));
    }

    center_scale_kernel<<<blocks, THREADS, 0, stream>>>(
        x.data_ptr<float>(), means.data_ptr<float>(),
        scales.data_ptr<float>(), output.data_ptr<float>(),
        channels, spatial, parts);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_bn_center = load_inline(
    name="convtranspose3d_bn_spatial_center",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["bn_center_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class _BatchNormCenter(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias, running_mean, running_var,
                batches, momentum, eps, training):
        output = _bn_center.bn_center_cuda(
            x, weight, running_mean, running_var,
            batches, momentum, eps, training,
        )
        ctx.training = training
        ctx.eps = eps
        if any(ctx.needs_input_grad[:3]):
            ctx.save_for_backward(
                x, weight, running_var.detach().clone()
            )
        return output

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

        centered_x = x - x.mean(spatial_dims, keepdim=True)
        centered_grad = grad_output - grad_output.mean(
            spatial_dims, keepdim=True
        )

        if ctx.training:
            variance, mean = torch.var_mean(
                x, dim=reduction_dims, unbiased=False, keepdim=True
            )
            invstd = torch.rsqrt(variance + ctx.eps)
            normalized = (x - mean) * invstd
            correction = (centered_grad * normalized).mean(
                reduction_dims, keepdim=True
            )
            grad_x = (
                centered_grad
                - centered_grad.mean(reduction_dims, keepdim=True)
                - normalized * correction
            ) * invstd * weight.view(1, -1, 1, 1, 1)
        else:
            invstd = torch.rsqrt(
                running_var.view(1, -1, 1, 1, 1) + ctx.eps
            )
            grad_x = (
                centered_grad * invstd * weight.view(1, -1, 1, 1, 1)
            )

        grad_weight = (grad_output * centered_x * invstd).sum(
            reduction_dims
        )
        grad_bias = torch.zeros_like(weight)
        return (
            grad_x, grad_weight, grad_bias,
            None, 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):
        x = self.conv_transpose(x)
        bn = self.batch_norm

        if (
            not x.is_cuda
            or x.dtype != torch.float32
            or not x.is_contiguous()
            or not bn.affine
            or not bn.track_running_stats
            or bn.momentum is None
            or bn.weight.dtype != torch.float32
            or x.numel() == 0
            or (bn.training and x.numel() // x.shape[1] <= 1)
        ):
            x = bn(x)
            return x - x.mean(dim=(2, 3, 4), keepdim=True)

        return _BatchNormCenter.apply(
            x, bn.weight, bn.bias, bn.running_mean, bn.running_var,
            bn.num_batches_tracked, bn.momentum, bn.eps, bn.training,
        )
```