load ae on gpu

app.py

@@ -21,6 +21,7 @@ def encode(init_image, torch_device, ae):
     init_image = torch.from_numpy(init_image).permute(2, 0, 1).float() / 127.5 - 1
     init_image = init_image.unsqueeze(0) 
     init_image = init_image.to(torch_device)
+    ae.to(torch_device)
     init_image = ae.encode(init_image.to()).to(torch.bfloat16)
     return init_image
 from torchvision import transforms