Commit 2d22d72c authored by flamelaw's avatar flamelaw

fix random sampling with pin_memory

parent a4a5735d
...@@ -134,7 +134,7 @@ class PersonalizedBase(Dataset): ...@@ -134,7 +134,7 @@ class PersonalizedBase(Dataset):
if self.tag_drop_out != 0 or self.shuffle_tags: if self.tag_drop_out != 0 or self.shuffle_tags:
entry.cond_text = self.create_text(entry.filename_text) entry.cond_text = self.create_text(entry.filename_text)
if self.latent_sampling_method == "random": if self.latent_sampling_method == "random":
entry.latent_sample = shared.sd_model.get_first_stage_encoding(entry.latent_dist) entry.latent_sample = shared.sd_model.get_first_stage_encoding(entry.latent_dist).to(devices.cpu)
return entry return entry
class PersonalizedDataLoader(DataLoader): class PersonalizedDataLoader(DataLoader):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment