Ожидалось, что все тензоры будут находиться на одном устройстве, но обнаружено как минимум два устройства: cpu и cuda:0, ⇐ Python

[Anonymous]

Я работаю над моделью, предназначенной для работы исключительно на графическом процессоре или на центральном процессоре, но поскольку я работаю над кластером HPC, я использую оба варианта.
Я' Я пытался преобразовать код для работы в моей настройке, но продолжал получать сообщение об ошибке. Ожидается, что все тензоры будут на одном устройстве.
Код выполняется в библиотеке Lightning и в исходном коде. установка имеет следующие параметры:
slurm_args = dict()
if os.environ.get("SLURM_NODELIST") is not None:
# Add SLURM arguments for distributed training
slurm_args = {
"accelerator": "gpu",
"devices": int(os.environ["SLURM_GPUS_ON_NODE"]),
"num_nodes": int(os.environ["SLURM_NNODES"]),
"strategy": "ddp",
}

...
trainer_args = {
"max_epochs": args.nb_epochs,
"log_every_n_steps": LOGGING_STEPS,
"val_check_interval": VAL_CHECK_STEPS,
"logger": wandb_logger,
"callbacks": callbacks,
"accelerator": accelerator,
**slurm_args,
}

...
trainer = lightning.Trainer(**trainer_args)

...
trainer.fit(**train_args, datamodule=datamodule)

Команда model.to("cuda") отсутствует нигде в коде и при добавлении ничего не меняет.
Кто-нибудь знает, какие настройки мне нужно изменить и на какие, чтобы можно было запускать под этими параметрами:
#SBATCH --time=60:00:00
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --mem=180000
#SBATCH --cpus-per-task=20
#SBATCH --gres=gpu:1


Подробнее здесь: https://stackoverflow.com/questions/790 ... -devices-c