Update app.py

app.py

@@ -94,7 +94,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.to(device).to(torch.bfloat16)
+    return vae.to(device=device, dtype=torch.bfloat16)
 
 def load_unet(unet_dir):
     unet_ckpt_path = unet_dir / "unet_diffusion_pytorch_model.safetensors"
@@ -103,7 +103,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).to(torch.bfloat16)
+    return transformer.to(device=device, dtype=torch.bfloat16)
 
 def load_scheduler(scheduler_dir):
     scheduler_config_path = scheduler_dir / "scheduler_config.json"