Fix PyTorch backend device placement so GNM runs on CUDA

gnm/shape/gnm_common.py

@@ -22,6 +22,22 @@
 _EPSILON = 1e-8
 
 
+def _constant_like(
+    values: Any,
+    reference_array: enpt.FloatArray['...'],
+    xnp: Any,
+    dtype: Any,
+) -> enpt.FloatArray['...']:
+  """Creates a constant array on the same device as `reference_array`.
+
+  For PyTorch this places the constant on `reference_array.device`; for other
+  backends it falls back to a plain `asarray`.
+  """
+  if enp.lazy.is_torch(reference_array):
+    return xnp.asarray(values, dtype=dtype, device=reference_array.device)
+  return xnp.asarray(values, dtype=dtype)
+
+
 def take(
     array: enpt.FloatArray['...'],
     indices: enpt.IntArray['...'],
@@ -48,6 +64,8 @@ def take(
     xnp = enp.lazy.get_xnp(array)
   if enp.lazy.is_torch_xnp(xnp):
     axis = axis % array.ndim
+    # index_select requires indices to live on the same device as `array`.
+    indices = xnp.as_tensor(indices, device=array.device)
     indices_shape = indices.shape
     if len(indices_shape) > 1:
       indices_flat = indices.reshape(-1)
@@ -184,7 +202,10 @@ def axis_angle_to_rotation_matrix(
 
   norm_squared = xnp.sum(xnp.square(axis_angle), axis=-1, keepdims=True)
   angle = xnp.sqrt(
-      xnp.maximum(norm_squared, xnp.asarray(epsilon, dtype=norm_squared.dtype))
+      xnp.maximum(
+          norm_squared,
+          _constant_like(epsilon, axis_angle, xnp, norm_squared.dtype),
+      )
   )
   axis = axis_angle / angle
 
@@ -275,8 +296,8 @@ def joint_transforms_world(
       joints, 2, (num_joints, 1, 4), dtype=joints.dtype
   )
   # Set the last element to 1.0.
-  bottom_row = bottom_row + xnp.asarray(
-      [0.0, 0.0, 0.0, 1.0], dtype=joints.dtype
+  bottom_row = bottom_row + _constant_like(
+      [0.0, 0.0, 0.0, 1.0], joints, xnp, joints.dtype
   )
 
   local_transforms = xnp.concatenate(

gnm/shape/gnm_pytorch.py

@@ -131,6 +131,17 @@ def _from_model_data(
     """Creates a PyTorch GNM instance from model data."""
     return cls._from_model_data_with_xnp(model_data, xnp=torch)
 
+  @classmethod
+  def from_gnm(cls, gnm: gnm_xnp.GNM) -> GNM:
+    """Creates a PyTorch GNM from another GNM, preserving its device."""
+    new_gnm = super().from_gnm(gnm)
+    # `from_gnm` reconstructs the model from a NumPy data dict, so the new model
+    # is created on CPU. If the source model lives on another device, move the
+    # reconstructed model there to match.
+    if isinstance(gnm, torch.nn.Module):
+      new_gnm = new_gnm.to(gnm.template_vertex_positions.device)
+    return new_gnm
+
   def compute_vertex_normals(
       self,
       vertices: torch.Tensor,

gnm/shape/gnm_pytorch_test.py

@@ -279,6 +279,7 @@ def test_prune_vertices(self, version: str, variant: Any):
         gnm_pytorch.GNMMajorVersion(version.removeprefix('v')),
         gnm_pytorch.GNMVariant(variant_str),
     )
+    gnm_pruned.to(self.device)
 
     keep_vertices = gnm_np.quads[0]
     gnm_pruned.prune_vertices(keep_vertices)

gnm/shape/gnm_xnp.py

@@ -379,6 +379,10 @@ def vertices_and_landmarks(
         translation=translation,
     )
     weights = enp.compat.astype(weights, vertices.dtype)
+    # The landmark tensors are cached lazily and may live on a different device
+    # than the (possibly moved) model; align them with the output vertices.
+    if enp.lazy.is_torch(vertices):
+      weights = weights.to(vertices.device)
 
     face_vertices = gnm_common.take(vertices, indices, axis=-2, xnp=self.xnp)
     landmarks = self.xnp.sum(face_vertices * weights[..., None], axis=-2)
@@ -634,6 +638,10 @@ def prune_vertices(self, keep_vertices: enpt.IntArray['V_pruned']) -> None:
     xnp = self.xnp
     num_vertices = self.num_vertices
     keep_vertices = xnp.asarray(keep_vertices, dtype=xnp.int32)
+    # Keep the index tensor on the same device as the model buffers so that all
+    # downstream gathers stay device-consistent.
+    if enp.lazy.is_torch(self.template_vertex_positions):
+      keep_vertices = keep_vertices.to(self.template_vertex_positions.device)
 
     self.template_vertex_positions = gnm_common.take(
         self.template_vertex_positions, keep_vertices, axis=0, xnp=xnp
@@ -729,8 +737,11 @@ def _scatter_indices(keep_vertices, num_vertices, xnp) -> Any:
         jnp.arange(len(keep_vertices), dtype=jnp.int32)
     )
   elif enp.lazy.is_torch_xnp(xnp):
-    mapper = xnp.full((num_vertices,), -1, dtype=xnp.int32)
-    mapper[keep_vertices] = xnp.arange(len(keep_vertices), dtype=xnp.int32)
+    device = keep_vertices.device
+    mapper = xnp.full((num_vertices,), -1, dtype=xnp.int32, device=device)
+    mapper[keep_vertices] = xnp.arange(
+        len(keep_vertices), dtype=xnp.int32, device=device
+    )
     return mapper
   else:
     mapper = np.full((num_vertices,), -1, dtype=np.int32)