ZeroGPU

demo.py

@@ -498,7 +498,7 @@ def generate_animation(
 
     print("Animation generated")
 
-    return synthetic_video.detach()  # B x C x T x H x W
+    return synthetic_video.detach().cpu()  # B x C x T x H x W
 
 
 @spaces.GPU
@@ -510,7 +510,8 @@ def decode_animation(latent_animation):
 
     # Convert to torch tensor if needed
     if not isinstance(latent_animation, torch.Tensor):
-        latent_animation = torch.from_numpy(latent_animation).to(device, dtype=dtype)
+        latent_animation = torch.from_numpy(latent_animation)
+    latent_animation = latent_animation.to(device, dtype=dtype)
 
     # Ensure shape is B x C x T x H x W
     if len(latent_animation.shape) == 4:  # [T, C, H, W]