Hi @Naplessss , While using this demo, in Block `Infer from ckpt`, a `KeyError: 'unimol'` was raised.

Hi @Naplessss , While using this demo, in Block Infer from ckpt, a KeyError: 'unimol' was raised.

  File "/opt/anaconda3/envs/unimol/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/opt/anaconda3/envs/unimol/lib/python3.9/site-packages/unicore/distributed/utils.py", line 164, in distributed_main
    main(args, **kwargs)
  File "/root/Uni-Mol/unimol/infer.py", line 47, in main
    task = tasks.setup_task(args)
  File "/opt/anaconda3/envs/unimol/lib/python3.9/site-packages/unicore/tasks/__init__.py", line 22, in setup_task
    return TASK_REGISTRY[args.task].setup_task(args, **kwargs)
KeyError: 'unimol'

This is the whole stack.

And I also print the TASK_REGISTRY in unicore/tasks/__init__.py. well, during running, that variable is {}. Could you please tell me what happens, and how to solve this problem?

_Originally posted by @zhoubay in https://github.com/dptech-corp/Uni-Mol/issues/56#issuecomment-1287846430_

[yuanbatou]

I set CUDA_VISIBLE_DEVICE to one GPU and can run this example. Hope it helps.

[zhoubay]

I set CUDA_VISIBLE_DEVICE to one GPU and can run this example. Hope it helps.

It works when I add a line in infer.py: os.environ['CUDA_VISIBLE_DEVICES']="0"

Thank you very much!