File size: 717 Bytes
4fd79a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 | from pathlib import Path
import torch
from spectral_world_models.data import DatasetConfig, generate_benchmark_npz, TransitionDataset
from spectral_world_models.models import build_model
def test_swm_forward(tmp_path: Path):
data = tmp_path / "tiny.npz"
generate_benchmark_npz(data, DatasetConfig(train_sequences=2, val_sequences=1, test_sequences=1, seq_len=4))
ds = TransitionDataset(data, "train")
batch = ds[0]
model = build_model("swm")
out = model(batch["image"].unsqueeze(0), batch["tokens"].unsqueeze(0), batch["action"].unsqueeze(0))
assert out["image_pred"].shape == (1, 1, 32, 32)
assert out["text_logits"].shape[1] == 5
assert torch.isfinite(out["image_pred"]).all()
|