load unet in bfloat16

app.py

@@ -45,7 +45,7 @@ def load_vae(vae_dir):
     vae = CausalVideoAutoencoder.from_config(vae_config)
     vae_state_dict = safetensors.torch.load_file(vae_ckpt_path)
     vae.load_state_dict(vae_state_dict)
-    return vae.cuda().to(torch.bfloat16)
+    return vae.to(device=device, dtype=torch.bfloat16)
 
 
 def load_unet(unet_dir):
@@ -55,7 +55,7 @@ def load_unet(unet_dir):
     transformer = Transformer3DModel.from_config(transformer_config)
     unet_state_dict = safetensors.torch.load_file(unet_ckpt_path)
     transformer.load_state_dict(unet_state_dict, strict=True)
-    return transformer.to(device)
+    return transformer.to(device=device, dtype=torch.bfloat16)
 
 
 def load_scheduler(scheduler_dir):