Столкнулся с такой же проблемой, почитал интернет и понял одну штуку. Да, хочется сразу всё загрузить в память и данные, и модель сети пожирнее и хотелось, чтобы torch сам управлял памятью/порциями данных и т.д. Фиг вам - это делаете вы сами своими ручками.
Если данных много и они не помещаются в память, то вы подгружаете их частями из файла. Так же и с видеопамятью - если туда всё не влазит, то это ваша проблема по разделению всего этого на минимально приемлемые части и подгрузка/выгрузка из памяти видеокарты. Если ваша модель НН очень большая, то вы должны своими ручками разделить её на части и обрабатывать её послойно загружай/выгружая в память gpu. Ну и всё остальное в таком же духе.
Короче, torch сам менеджментом памяти на gpu не занимается. Если данные в память видеокарты не влазят, то он просто выпадает в ошибку.
Есть какие-то библиотеки, которые работаю обёрткой для torch и занимаются этим менеджментом, но я в ту сторону ещё не копал.
Если у кого-то есть другая информация, то с радостью почитаю.
P.S. В частности про Stable Diffusion. Там есть такая настройка как "Low memory mode (for 8Gb GPUs)", что можно интерпретировать как "режим для отстойных видях, у которых всего 8Гб памяти" ... а вы пытаетесь запустить на 6Гб ... сомневаюсь.
Сложный
Сложный
