Update app.py

app.py

@@ -19,7 +19,7 @@ dtype = torch.float16
 pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
 ).to("cuda")
-pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16)
+pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
 torch.cuda.empty_cache()
 
 # Inference function