import torch, json, numpy as np
from model_def import load_model

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

with open("config.json") as f:
    cfg = json.load(f)

model = load_model("pytorch_model.bin", "config.json", device=DEVICE)

x = np.random.randn(1, cfg["input_size"]).astype("float32")
x_t = torch.from_numpy(x).to(DEVICE)

with torch.no_grad():
    y_hat = model(x_t).cpu().numpy()

print("Pred shape:", y_hat.shape)
print("Pred sample:", y_hat[0][: min(5, y_hat.shape[1])])