george2cool36's picture
Initial upload: model weights, config, metrics, README, model_def.py, inference.py
60fd15e verified
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])])