DataParallel multi-gpu RuntimeError: чанк ожидает как минимум одномерный тензор ⇐ Python

[Anonymous]

Я пытаюсь запустить свою модель на нескольких графических процессорах с помощью DataParallel, установив model = nn.DataParallel(model).cuda(), но каждый раз получаю эту ошибку:

RuntimeError: чанк ожидает как минимум одномерный тензор (чанк по адресу
/pytorch/aten/src/ATen/native/TensorShape.cpp:184).

Мой код правильный. Кто-нибудь знает, что не так?
Я попробовал установить параметр device_ids=[0,1], а также CUDA_VISIBLE_DEVICES на терминале. Также пробовал разные размеры партий.

Подробнее здесь: https://stackoverflow.com/questions/561 ... ional-tens