1 files changed, 5 insertions, 6 deletions

diff --git a/training/functional.py b/training/functional.py
index 56c2995..fd3f9f4 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -68,14 +68,13 @@ class TrainingStrategy():
     prepare: TrainingStrategyPrepareCallable
 
 
-def get_models(pretrained_model_name_or_path: str):
+def get_models(pretrained_model_name_or_path: str, torch_dtype=torch.float32):
     tokenizer = MultiCLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer')
-    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder')
-    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder='vae')
-    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder='unet')
+    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder', torch_dtype=torch_dtype)
+    vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder='vae', torch_dtype=torch_dtype)
+    unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder='unet', torch_dtype=torch_dtype)
     noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder='scheduler')
-    sample_scheduler = UniPCMultistepScheduler.from_pretrained(
-        pretrained_model_name_or_path, subfolder='scheduler')
+    sample_scheduler = UniPCMultistepScheduler.from_pretrained(pretrained_model_name_or_path, subfolder='scheduler')
 
     return tokenizer, text_encoder, vae, unet, noise_scheduler, sample_scheduler