device — где лежит память (cpu, cuda:0, mps); перенос делают .to(device) / .cuda() / .cpu().
Разбор
- Сначала проверяют
torch.cuda.is_available().
- Модель и батч должны быть на одном device.
- Частые
.cpu()/.cuda() в цикле убивают скорость.
- На Apple Silicon часто
mps, с оговорками по операциям.
Пример
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
batch = batch.to(device)
Итог
Единый device для модели и данных — базовое правило обучения.