From e54b6311ef7942af0b7b10c49279ce1e71c4c145 Mon Sep 17 00:00:00 2001 From: Kohya S Date: Sun, 7 May 2023 10:23:25 +0900 Subject: [PATCH] do not save cuda_rng_state if no cuda closes #390 --- library/train_util.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 7b553363..ea433979 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -3344,7 +3344,7 @@ def sample_images( os.makedirs(save_dir, exist_ok=True) rng_state = torch.get_rng_state() - cuda_rng_state = torch.cuda.get_rng_state() + cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None with torch.no_grad(): with accelerator.autocast(): @@ -3451,7 +3451,8 @@ def sample_images( torch.cuda.empty_cache() torch.set_rng_state(rng_state) - torch.cuda.set_rng_state(cuda_rng_state) + if cuda_rng_state is not None: + torch.cuda.set_rng_state(cuda_rng_state) vae.to(org_vae_device)