Skip to content

Fix PyTorch backend device placement so GNM runs on CUDA#41

Open
AvirupJU wants to merge 1 commit into
google:mainfrom
AvirupJU:main
Open

Fix PyTorch backend device placement so GNM runs on CUDA#41
AvirupJU wants to merge 1 commit into
google:mainfrom
AvirupJU:main

Conversation

@AvirupJU

Copy link
Copy Markdown

Summary

The PyTorch backend is unusable on GPU: after moving a model to CUDA
(e.g. gnm.cuda()), evaluation fails with device-mismatch errors such as:

RuntimeError: Expected all tensors to be on the same device, but got index is
on cpu, different from other tensors on cuda:0 (... wrapper_CUDA__index_select)

Several tensors are built on CPU (from Python literals, lazily-cached landmark
data, or freshly-reconstructed models) and then combined with model buffers that
have been moved to GPU. This never surfaces in CI because the CI matrix is
CPU-only (Linux/macOS/Windows, no GPU runner), so the whole GPU path is
untested.

With a CUDA device visible, gnm_pytorch_test.py fails 28 / 33; with
CUDA_VISIBLE_DEVICES="" all pass. This PR makes the PyTorch backend
device-consistent so the model, landmarks, pruning, and from_gnm all work on
GPU.

Changes

All changes are guarded by is_torch / is_torch_xnp checks, so the NumPy,
JAX, and TensorFlow backends are unaffected.

File Change
gnm/shape/gnm_common.py take() moves indices onto the array's device before index_select. New _constant_like() helper places literal constants on the reference tensor's device for torch; used in axis_angle_to_rotation_matrix (epsilon) and joint_transforms_world (the [0,0,0,1] homogeneous row).
gnm/shape/gnm_xnp.py Align lazily-cached landmark weights with the output device; place keep_vertices (prune_vertices) and the _scatter_indices mapper on the model's device.
gnm/shape/gnm_pytorch.py Override from_gnm to move the reconstructed model onto the source model's device (from_gnm rebuilds from a NumPy dict, which always lands on CPU).
gnm/shape/gnm_pytorch_test.py Fix a latent test bug: test_prune_vertices built gnm_pruned via from_local (CPU) but fed it GPU parameters. Added gnm_pruned.to(self.device).

Testing

Validated locally on CUDA (torch 2.13 + cu130, Python 3.13):

  • gnm_pytorch_test.py: 33 passed on GPU, 33 passed on CPU.
  • Shared / other-backend suites (gnm_numpy_test.py, gnm_xnp_test.py,
    gnm_base_test.py, gnm_utils_test.py, gnm_landmarks_test.py):
    155 passed, 2 skipped — no regressions.
cd gnm/shape
python -m pytest gnm_pytorch_test.py -q                       # GPU
CUDA_VISIBLE_DEVICES="" python -m pytest gnm_pytorch_test.py -q # CPU

@google-cla

google-cla Bot commented Jul 21, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant