Fix meta-device init crash on transformers>=5.0 (add device='cpu' to linspace)
#16
by eliem - opened
transformers>=5.0 initializes models on the meta device, where the stochastic-depth schedule built at init via [x.item() for x in torch.linspace(0, drop_path_rate, depth)] raises: RuntimeError: Tensor.item() cannot be called on meta tensors (reported in github.com/mahmoodlab/TITAN issue #43).
This PR pins the two linspace calls (conch_v1_5.py L336, vision_transformer.py L316) to device='cpu', which is immune to the meta-device context and produces bit-identical drop-path values on every transformers/torch version. No other init-time .item() calls exist in the remote code (the ones in vision_transformer.py L35-36 run at forward time on real tensors).