Failed to run DiffDock Inference on CPU

[Franreno]

Hi there!

I'm running DiffDock's inference on a VM that does not have a GPU or CUDA installed. While executing the model using the command below, it runs successfully up to a point but then crashes. The issue arises because the InferenceDataset class attempts to move a model to the GPU using model.eval().cuda() without first checking if CUDA is available.

Command

env TORCH_HOME=./ micromamba run -n diffdock python ./DiffDock/inference.py \
    --config ./default_inference_args.yaml \
    --protein_sequence GIQSYCTPPYSVLQDPPQPVV \
    --ligand "COc(cc1)ccc1C#N" \
    --samples_per_complex 1000

stderr

Traceback (most recent call last):
  File "/workdir/output/artifacts/./DiffDock/inference.py", line 318, in <module>
    main(_args)
  File "/workdir/output/artifacts/./DiffDock/inference.py", line 173, in main
    test_dataset = InferenceDataset(out_dir=args.out_dir, complex_names=complex_name_list, protein_files=protein_path_list,
  File "/workdir/output/artifacts/DiffDock/utils/inference_utils.py", line 169, in __init__
    model = model.eval().cuda()
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/nn/modules/module.py", line 749, in cuda
    return self._apply(lambda t: t.cuda(device))
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/nn/modules/module.py", line 641, in _apply
    module._apply(fn)
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/nn/modules/module.py", line 664, in _apply
    param_applied = fn(param)
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/nn/modules/module.py", line 749, in <lambda>
    return self._apply(lambda t: t.cuda(device))
  File "/home/appuser/micromamba/envs/diffdock/lib/python3.9/site-packages/torch/cuda/__init__.py", line 229, in _lazy_init
    torch._C._cuda_init()
RuntimeError: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx

---

I've noticed that in the same InferenceDataset class, there is a proper CUDA availability check when initializing the esm2_t33_650M_UR50D model. (reference).

            print("Generating ESM language model embeddings")
            model_location = "esm2_t33_650M_UR50D"
            model, alphabet = pretrained.load_model_and_alphabet(model_location)
            model.eval()
            if torch.cuda.is_available():
                model = model.cuda()

However, for the esmfold_v1 model, no such check is performed before calling.cuda() (reference).

            print("generating missing structures with ESMFold")
            model = esm.pretrained.esmfold_v1()
            model = model.eval().cuda()

Would it be possible to update this section to move the model to the GPU only if CUDA is available? This way, DiffDock can run more flexibly on CPU-only environments.

Thanks!