search

TencentAI4S / tfold

open source code for Tencent tFold
Other
57 stars 9 forks source link

'RuntimeError: input must be a CUDA tensor' running in Docker #4

Closed oliverfleetwood closed 2 months ago

oliverfleetwood commented 2 months ago

Hi,

I run the examples for the antigen-antibody complex in your README inside a docker container based on your image. I get the following error:

python projects/tfold_ag/predict.py --pid_fpath=examples/prot_ids.abag.txt --fas_dpath=examples/fasta.files --msa_fpath=examples/msa.files/8df5_R.a3m --pdb_dpath=examples/pdb.files.abag

output

[2024-04-15 20:28:05,534] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
WARNING:root:using DS4Sci_EvoformerAttention
[04/15 20:28:08 root]: restoring the pre-trained tFold-PPI-SeqPT model ...
[04/15 20:28:15 root]: model weights restored from params/esm_ppi_650m.pth
[04/15 20:28:15 root]: restoring the pre-trained AlphaFold2 model ...
[04/15 20:28:22 root]: model weights restored from params/alphafold_4_ptm.pth
[04/15 20:28:22 root]: restoring the pre-trained tFold-Ag model ...
model:
  arch: tfold_ag
  docking:
    ligand_c_s: 192
    ligand_c_z: 128
    num_2d_layers: 32
    num_2d_recycles: 1
    num_3d_layers: 8
    num_3d_recycles: 8
    num_recycles: 1
    receptor_c_s: 256
    receptor_c_z: 128
    use_icf: false
  ligand:
    c_s: 192
    c_z: 128
    init_c_s: 1280
    init_c_z: 660
    num_2d_layers: 16
    num_2d_recycles: 1
    num_3d_layers: 8
    num_3d_recycles: 8
    num_recycles: 1
    position_embedding_dim: 64
    tm_enabled: true
    use_residue_embedding: false
  plm:
    c_s: 1280
    c_z: 660

[04/15 20:28:24 root]: restore the pre-trained tFold-Ag model params/tfold_ag.pth
[04/15 20:28:24 root]: 8df5_R antigen feature predicted using pretrained structure prediction model
Traceback (most recent call last):
  File "projects/tfold_ag/predict.py", line 81, in <module>
    main()
  File "projects/tfold_ag/predict.py", line 73, in main
    predict(args.pid_fpath, args.fas_dpath,
  File "projects/tfold_ag/predict.py", line 67, in predict
    predictor.run_batch(pid_fpath, fas_dpath, msa_fpath, pdb_dpath, icf_dpath, chunk_size=None)
  File "/miniconda/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/root/tfold/./tfold/deploy/ag_predictor.py", line 133, in run_batch
    ag_feat, _ = self.psp_featurizer(msa_path)
  File "/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/root/tfold/./tfold/deploy/psp_featurizer.py", line 101, in forward
    result, ptm = self._forward_impl(process_feature_dict)
  File "/root/tfold/./tfold/deploy/psp_featurizer.py", line 131, in _forward_impl
    outputs, m_1_prev, z_prev, x_prev = self.model.iteration(feats, prevs, _recycle=(num_iters > 1))
  File "/miniconda/lib/python3.8/site-packages/openfold/model/model.py", line 366, in iteration
    z = self.extra_msa_stack(
  File "/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/miniconda/lib/python3.8/site-packages/openfold/model/evoformer.py", line 1007, in forward
    m, z = b(m, z)
  File "/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/miniconda/lib/python3.8/site-packages/openfold/model/evoformer.py", line 518, in forward
    self.msa_att_row(
  File "/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/miniconda/lib/python3.8/site-packages/openfold/model/msa.py", line 266, in forward
    m = self._chunk(
  File "/miniconda/lib/python3.8/site-packages/openfold/model/msa.py", line 121, in _chunk
    return chunk_layer(
  File "/miniconda/lib/python3.8/site-packages/openfold/utils/chunk_utils.py", line 299, in chunk_layer
    output_chunk = layer(**chunks)
  File "/miniconda/lib/python3.8/site-packages/openfold/model/msa.py", line 101, in fn
    return self.mha(
  File "/miniconda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/miniconda/lib/python3.8/site-packages/openfold/model/primitives.py", line 492, in forward
    o = attention_core(q, k, v, *((biases + [None] * 2)[:2]))
  File "/miniconda/lib/python3.8/site-packages/torch/autograd/function.py", line 506, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
  File "/miniconda/lib/python3.8/site-packages/openfold/utils/kernel/attention_core.py", line 47, in forward
    attn_core_inplace_cuda.forward_(
RuntimeError: input must be a CUDA tensor
Cloud-Rambler commented 2 months ago

Hello, this is because the docker container does not detect GPU, you can use this command to start the docker container：

docker run -it --gpus all ai4science/tfold:latest /bin/bash

Driver versions are preferred to support cuda11.7, which you can check with 

nvidia-smi
oliverfleetwood commented 2 months ago

worked with the right version of cuda, thanks!