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

Load training dataset tensors with weights_only

Browse Source
This commit is contained in:
CFionaBF 2026-06-03 08:22:08 -04:00
parent bd7da053ae
commit 43726d744b
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

@ -1565,7 +1565,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"])