From 94ee49b1612824366a8631ea069b2a1fa5c73720 Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Thu, 18 Jun 2026 12:30:57 -0700 Subject: [PATCH] harden: load training-dataset shards with weights_only=True (#14543) LoadTrainingDataset was the only torch.load call in the codebase without weights_only=True; comfy/utils.py and comfy/sd1_clip.py already pass it. Recent PyTorch defaults to weights_only=True, so this is defense-in-depth for installs pinned to older PyTorch. Verified a typical shard (latents + standard conditioning) round-trips cleanly under weights_only=True. --- comfy_extras/nodes_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 0253b4b4f..73fe75b7f 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -1583,7 +1583,7 @@ class LoadTrainingDataset(io.ComfyNode): shard_path = os.path.join(dataset_dir, shard_file) with open(shard_path, "rb") as f: - shard_data = torch.load(f) + shard_data = torch.load(f, weights_only=True) all_latents.extend(shard_data["latents"]) all_conditioning.extend(shard_data["conditioning"])