move to device on stateless gpu

app.py

@@ -70,9 +70,7 @@ def run_rerun(path_to_video):
         rr.log("world/camera/image", rr.Image(frame))
         yield stream.read()
 
-        image = transform(frame)
-        image = image.to(device)
-        depth, focal_length = estimate_depth(image)
+        depth, focal_length = estimate_depth(frame)
 
         rr.log(
             "world/camera",
@@ -96,7 +94,9 @@ def run_rerun(path_to_video):
 
 
 @spaces.GPU(duration=20)
-def estimate_depth(image):
+def estimate_depth(frame):
+    image = transform(frame)
+    image = image.to(device)
     prediction = model.infer(image)
     depth = prediction["depth"].squeeze().detach().cpu().numpy()
     focal_length = prediction["focallength_px"].item()