import svetoviz_webgpu as sv

# 1. Prepare 2D image data
input_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0)

# 2. Define BatchNorm2d
# num_features: Number of channels (e.g., 3 for RGB)
bn2d = nn.BatchNorm2d(num_features=3)

def terminal_callback(buffer, message, images, files):
    bn2d.train()
    output = bn2d(input_tensor)

    buffer.send_system_message(f"Image Resolution: {input_tensor.shape[2:]}")
    buffer.send_system_message(f"Channels Normalized: {bn2d.num_features}")
    buffer.send_system_message("Status: Each channel normalized across the spatial plane.")

# 3. Start the interactive session
sv.pytorch_web(module=bn2d, terminal_callback=terminal_callback)