Fix device allocation for top-down models

dlclive/pose_estimation_pytorch/runner.py

@@ -72,11 +72,33 @@ def update(self, pose: torch.Tensor, w: int, h: int) -> None:
         num_det, num_kpts = pose.shape[:2]
         size = max(w, h)
 
-        bboxes = torch.zeros((num_det, 4))
-        bboxes[:, :2] = torch.min(torch.nan_to_num(pose, size)[..., :2], dim=1)[0] - self.margin
-        bboxes[:, 2:4] = torch.max(torch.nan_to_num(pose, 0)[..., :2], dim=1)[0] + self.margin
-        bboxes = torch.clip(bboxes, min=torch.zeros(4), max=torch.tensor([w, h, w, h]))
-        self._detections = dict(boxes=bboxes, scores=torch.ones(num_det))
+        bboxes = pose.new_zeros((num_det, 4))
+
+        bboxes[:, :2] = (
+            torch.min(
+                torch.nan_to_num(pose[..., :2], nan=float(size)),
+                dim=1,
+            )[0]
+            - self.margin
+        )
+
+        bboxes[:, 2:4] = (
+            torch.max(
+                torch.nan_to_num(pose[..., :2], nan=0.0),
+                dim=1,
+            )[0]
+            + self.margin
+        )
+
+        minimum = pose.new_zeros(4)
+        maximum = pose.new_tensor([w, h, w, h])
+
+        bboxes = torch.clamp(bboxes, min=minimum, max=maximum)
+
+        self._detections = {
+            "boxes": bboxes,
+            "scores": pose.new_ones(num_det),
+        }
         self._age += 1
 
 
@@ -188,7 +210,7 @@ def get_pose(self, frame: np.ndarray) -> np.ndarray:
 
             frame_batch, offsets_and_scales = self._prepare_top_down(tensor, detections)
             if len(frame_batch) == 0:
-                offsets_and_scales = [(0, 0), 1]
+                offsets_and_scales = [((0, 0), (1.0, 1.0))]
             tensor = frame_batch  # still CHW, batched
 
         if self.dynamic is not None:
@@ -357,21 +379,24 @@ def _postprocess_top_down(
         """Post-processes pose for top-down models."""
         if len(batch_pose) == 0:
             bodyparts, coords = batch_pose.shape[-2:]
-            return torch.zeros((0, bodyparts, coords))
+            return batch_pose.new_zeros((0, bodyparts, coords))
 
         poses = []
+
         for pose, (offset, scale) in zip(batch_pose, offsets_and_scales, strict=False):
-            poses.append(
-                torch.cat(
-                    [
-                        pose[..., :2] * torch.tensor(scale) + torch.tensor(offset),
-                        pose[..., 2:3],
-                    ],
-                    dim=-1,
-                )
+            scale_tensor = pose.new_tensor(scale)
+            offset_tensor = pose.new_tensor(offset)
+
+            transformed_pose = torch.cat(
+                [
+                    pose[..., :2] * scale_tensor + offset_tensor,
+                    pose[..., 2:3],
+                ],
+                dim=-1,
             )
+            poses.append(transformed_pose)
 
-        return torch.cat(poses)
+        return torch.stack(poses, dim=0)
 
 
 def _parse_device(device: str | None) -> str: