import svetoviz_webgpu as sv

# 1. Prepare 1D Signal (Batch=1, Channels=4, Length=8)
input_tensor = torch.ones(1, 4, 8)

# 2. Define Dropout1d (Spatial Dropout)
dropout = nn.Dropout1d(p=0.5)

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

    # Check which entire channels are zeroed
    channel_sums = output.sum(dim=2).squeeze()
    active_channels = (channel_sums > 0).sum().item()

    buffer.send_system_message(f"Active Channels: {active_channels} / 4")
    buffer.send_system_message("Logic: Entire 1D feature maps are dropped, not just individual pixels.")

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