Ошибка «CUDA error: an illegal memory access was encountered»

Возможные варианты решения:

Запустить скрипт обучения с параметром CUDA_LAUNCH_BLOCKING=1 для получения трассировки ошибки.
Уменьшить batch size и убедиться, что никакие другие процессы (например, Jupyter) не занимают видеопамять.
Использовать образ с более новыми версиями CUDA и PyTorch (см. Образы для Jupyter Server).
Применить contiguous() к тензорам, из-за которых возникает ошибка.

Подробнее о решении проблемы на GitHub.