Loading a model saved on CUDA fails on CPU-only machines

[robosimon]

Summary

CEBRA.load(...) fails with RuntimeError: No CUDA GPUs are available when loading a checkpoint that was saved on a CUDA device from an environment where no GPU is available (e.g. CPU-only machine or notebook, or PyTorch built without CUDA).

Environment

• OS: macOS 15.2 (Darwin 25.2.0, arm64)
• CEBRA version: 0.6.0
• PyTorch version: 2.10.0
• Device: CPU-only (CUDA not available)

Steps to Reproduce

1. Train and save a CEBRA model on a machine with CUDA:

   import cebra
   import numpy as np
   
   X = np.random.uniform(0, 1, (100, 5))
   model = cebra.CEBRA(max_iterations=10, device='cuda').fit(X)
   model.save('checkpoint.pt')

2. On a CPU-only machine (or with CUDA unavailable), load the model:

   from cebra import CEBRA
   model = CEBRA.load('checkpoint.pt')

3. Observe error:

   RuntimeError: No CUDA GPUs are available
   

Root Cause

• The checkpoint stores the training device in state['device_'] (e.g., 'cuda' or 'cuda:0')
• _load_cebra_with_sklearn_backend() calls .to(state['device_']) on model, criterion, and solver
• This triggers PyTorch CUDA initialization; if no GPU available, it raises the error
• map_location=torch.device('cpu') in torch.load() only affects tensor loading, not subsequent .to() calls

Expected Behavior

Loading should succeed on CPU-only environments: the model should gracefully fall back to CPU when CUDA is not available, regardless of the device used when saving.

Related

This fix is intended to be upstreamed to AdaptiveMotorControlLab/CEBRA.

[robosimon]

Fix Available

A fix for this issue has been implemented in PR #2.

Summary of the Fix

The fix adds a _resolve_checkpoint_device() helper function that detects when a checkpoint was saved on CUDA but CUDA is not available at load time, and automatically falls back to CPU.

Changes

• cebra/integrations/sklearn/cebra.py: Added device resolution logic
• tests/test_sklearn.py: Added comprehensive tests that verify the fix

Test Coverage

• 8 parametrized test cases covering various CUDA device formats
• Tests use monkeypatch to mock CPU-only environments
• All tests pass ✓

Please review the PR when ready.