Training ORGAN on cuda error

[Binbose]

Hey,

when I train ORGAN via python scripts/run.py --model organ --train_path ./my_organ/data/hce_smiles.smi --checkpoint_dir ./my_orgapt --device cuda:0 --metrics ./my_organ/metrics it throws the error

Traceback (most recent call last):
  File "scripts/run.py", line 219, in <module>
    main(config)
  File "scripts/run.py", line 199, in main
    train_model(config, model, train_path, test_path)
  File "scripts/run.py", line 127, in train_model
    trainer_script.main(model, trainer_config)
  File "/home/luca/Projects/benchmark_level_up/local/moses/scripts/train.py", line 62, in main
    trainer.fit(model, train_data, val_data)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/molsets-1.0-py3.6.egg/moses/organ/trainer.py", line 370, in fit
    gen_val_loader, logger)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/molsets-1.0-py3.6.egg/moses/organ/trainer.py", line 104, in _pretrain_generator
    criterion, optimizer)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/molsets-1.0-py3.6.egg/moses/organ/trainer.py", line 70, in _pretrain_generator_epoch
    outputs, _, _ = model.generator_forward(prevs, lens)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/molsets-1.0-py3.6.egg/moses/organ/model.py", line 94, in generator_forward
    return self.generator(*args, **kwargs)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/torch-1.7.1-py3.6-linux-x86_64.egg/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/molsets-1.0-py3.6.egg/moses/organ/model.py", line 23, in forward
    x = pack_padded_sequence(x, lengths, batch_first=True)
  File "/home/luca/anaconda3/envs/moses/lib/python3.6/site-packages/torch-1.7.1-py3.6-linux-x86_64.egg/torch/nn/utils/rnn.py", line 244, in pack_padded_sequence
    _VF._pack_padded_sequence(input, lengths, batch_first)
RuntimeError: 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor

Seems like changing line 23 in moses/organ/model.py from x = pack_padded_sequence(x, lengths.cpu(), batch_first=True) to x = pack_padded_sequence(x, lengths.cpu(), batch_first=True) fixes the problem though