enable zerogpu

app.py

@@ -212,12 +212,14 @@ if NEW_MODEL:
     # ckpt_state_dict = torch.load(model_path)['model_state_dict']
     ckpt_state_dict = torch.load(model_path, map_location='cpu')['ema_state_dict']
     missing_keys, extra_keys = model.load_state_dict(ckpt_state_dict, strict=False)
+    model = model.to(device)
     model.eval()
     print(missing_keys, extra_keys)
     assert len(missing_keys) == 0
     vae_state_dict = torch.load(vae_path)['state_dict']
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False) # .to(device)
+    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False)
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
+    autoencoer = autoencoder.to(device)
     autoencoder.eval()
     assert len(missing_keys) == 0
 else: