File size: 499 Bytes
60fd15e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
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])])
|