# /// script
# requires-python = ">=3.10"
# dependencies = ["wildedge-sdk", "torch", "numpy"]
#
# [tool.uv.sources]
# wildedge-sdk = { path = "..", editable = true }
# ///
"""PyTorch integration example. Run with: uv run pytorch_example.py"""

import torch
import torch.nn as nn

import wildedge

client = wildedge.WildEdge(
    app_version="1.0.0",  # uses WILDEDGE_DSN if set; otherwise no-op
)


class SimpleClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 10),
        )

    def forward(self, x):
        return self.net(x)


model = client.load(SimpleClassifier)
device = next(model.parameters()).device
batch = torch.randn(8, 128, device=device)

model.eval()

with torch.inference_mode():
    for _ in range(3):
        output = model(batch)
        print("output shape:", output.shape)
