Inference shim: checkpoint fails to load under PyTorch >= 2.6 (weights_only default rejects omegaconf globals)

[isayev]

Summary

scripts/loqi_inference_shim.py cannot load data/loqi.ckpt under PyTorch ≥ 2.6. The Graph3DInterpolantModel.load_from_checkpoint(...) call (shim line 61) raises UnpicklingError: the checkpoint's hyper-parameters contain omegaconf.DictConfig globals, and PyTorch 2.6 changed torch.load's default to weights_only=True, which refuses them. Lightning's internal load inside load_from_checkpoint does not pass weights_only=False.

The shim's own peek load (line 55) already passes weights_only=False, so only the load_from_checkpoint path is affected — easy to miss.

Traceback (abridged)

File ".../scripts/loqi_inference_shim.py", line 61, in _load_model
    model = Graph3DInterpolantModel.load_from_checkpoint(str(ckpt_path), ...)
...
File ".../torch/serialization.py", line 1529, in load
    raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_pickle.UnpicklingError: Weights only load failed ...
    WeightsUnpickler error: Unsupported global: GLOBAL omegaconf.dictconfig.DictConfig
    was not an allowed global by default.

Workaround

Run with TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 — Lightning honors it and forces weights_only=False, and the shim loads and samples normally.

Suggested fix

The checkpoint is first-party/trusted, so either:

• allow-list the omegaconf globals before loading:

  import omegaconf, torch
  torch.serialization.add_safe_globals([
      omegaconf.dictconfig.DictConfig,
      omegaconf.listconfig.ListConfig,
      omegaconf.base.ContainerMetadata,
      omegaconf.nodes.AnyNode,
  ])

• or pass weights_only=False through the checkpoint load in _load_model.

Either makes the shim work on stock PyTorch ≥ 2.6 without the env var.

Environment

PyTorch 2.9.1+cu128, Lightning 2.6.1, LoQI main @ 3f1c3fe, data/loqi.ckpt.