import torch


state = torch.load("/artifacts/fl2va-sampler-reference/initial.pt", map_location="cpu", weights_only=False)
for name, value in state.items():
    print(name, type(value), getattr(value, "shape", None))
    if hasattr(value, "unbind"):
        try:
            print("unbind", [(type(item), getattr(item, "shape", None)) for item in value.unbind()])
        except Exception as error:
            print("unbind error", error)