From 1e2c20ca3e6601ee2a64d8498e8ed957dfff493e Mon Sep 17 00:00:00 2001 From: Cory Stephenson Date: Fri, 4 Oct 2024 06:05:09 +0000 Subject: [PATCH] Update to new torch autocast --- diffusion/models/precomputed_text_latent_diffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diffusion/models/precomputed_text_latent_diffusion.py b/diffusion/models/precomputed_text_latent_diffusion.py index d1ee9136..31acdb2f 100644 --- a/diffusion/models/precomputed_text_latent_diffusion.py +++ b/diffusion/models/precomputed_text_latent_diffusion.py @@ -189,7 +189,7 @@ def set_rng_generator(self, rng_generator: torch.Generator): self.rng_generator = rng_generator def encode_images(self, inputs, dtype=torch.bfloat16): - with torch.amp.autocast('cuda', enabled=False): + with torch.autocast(device_type='cuda', enabled=False): latents = self.vae.encode(inputs.to(dtype))['latent_dist'].sample().data latents = (latents - self.latent_mean) / self.latent_std # scale latents return latents