File size: 576 Bytes
626ab93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | from __future__ import annotations
import json
from model import TuringSurrogate
from safetensors.torch import load_file
from train import ARTIFACT_DIR, rollout, save_comparison_gif
def main() -> None:
model = TuringSurrogate()
model.load_state_dict(load_file(ARTIFACT_DIR / "model.safetensors"))
metrics, physics_frames, neural_frames = rollout(model)
save_comparison_gif(
physics_frames,
neural_frames,
ARTIFACT_DIR / "physics_vs_neural.gif",
)
print(json.dumps(metrics, indent=2))
if __name__ == "__main__":
main()
|