| 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() | |