wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-24 00:39:30 +08:00
Code Issues Actions 18 Packages Projects Releases Wiki Activity

Load training dataset tensors with weights_only

Browse Source
This commit is contained in:
CFionaBF 2026-06-07 11:20:20 -04:00
parent 2cdaaf4a25
commit 1b877e583f
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy_extras/nodes_dataset.py
View File

@ -1583,7 +1583,7 @@ class LoadTrainingDataset(io.ComfyNode):
shard_path = os.path.join(dataset_dir, shard_file) shard_path = os.path.join(dataset_dir, shard_file)
with open(shard_path, "rb") as f: 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_latents.extend(shard_data["latents"])
all_conditioning.extend(shard_data["conditioning"]) all_conditioning.extend(shard_data["conditioning"])