tests/test_hdri_encoder_torch_load.py

from __future__ import annotations

import unittest
from pathlib import Path


HDRI_PATH = Path(__file__).resolve().parents[1] / "trellis" / "models" / "structured_latent_vae" / "hdri_encoder.py"


class HdriEncoderTorchLoadTests(unittest.TestCase):
    def test_load_weights_uses_full_unpickle_on_cpu(self) -> None:
        text = HDRI_PATH.read_text(encoding="utf-8")
        self.assertIn("weights_only=False", text)
        self.assertIn("torch.device(\"cpu\")", text)


if __name__ == "__main__":
    unittest.main()