Relaxation Error Running Locally

Expected Behavior

I am unable to use --amber to relax models.

Current Behavior

First model will fold but crash immediately when the relaxation step begins.

Steps to Reproduce (for bugs)

conda update conda
conda create -n colabfold python=3.7
conda activate colabfold
pip install --upgrade pip
pip install "colabfold[alphafold] @ git+https://github.com/sokrypton/ColabFold"
conda install cudatoolkit
conda install cudnn
pip install "jax[cuda11_cudnn805]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
conda install -y -c conda-forge -c bioconda kalign2=2.04 hhsuite=3.3.0
conda install -y -c conda-forge openmm=7.5.1 pdbfixer (This does cause a small error that says failed with initial frozen solve but does complete, maybe this is where the problem lies)
colabfold_batch 6CDX.fasta out1/ --num-models 1 --amber --use-gpu-relax

ColabFold Output (for bugs)

2022-06-27 20:23:01,839 Running colabfold 1.3.0 (2a47c6f1459fbbdb5242cbc62173f9b513813cfa)
Downloading alphafold2 weights to /home/ubuntu/.cache/colabfold: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3.47G/3.47G [00:13<00:00, 277MB/s]
WARNING: You are welcome to use the default MSA server, however keep in mind that it's a limited shared resource only capable of processing a few thousand MSAs per day. Please submit jobs only from a single IP address. We reserve the right to limit access to the server case-by-case when usage exceeds fair use.

If you require more MSAs:

• You can precompute all MSAs with colabfold_search or

• You can host your own API and pass it to --host-url
  2022-06-27 20:23:16,563 generated new fontManager
  2022-06-27 20:23:18,610 Found 6 citations for tools or databases
  2022-06-27 20:23:22,591 Query 1/1: 6CDX (length 36)
  COMPLETE: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 150/150 [elapsed: 00:01 remaining: 00:00]
  2022-06-27 20:23:23,763 Running model_3
  2022-06-27 20:23:46,034 model_3 took 19.9s (3 recycles) with pLDDT 74.4 and ptmscore 0.432
  Traceback (most recent call last):
  File "/home/ubuntu/.local/bin/colabfold_batch", line 8, in
  sys.exit(main())
  File "/home/ubuntu/.local/lib/python3.8/site-packages/colabfold/batch.py", line 1724, in main
  run(
  File "/home/ubuntu/.local/lib/python3.8/site-packages/colabfold/batch.py", line 1369, in run
  outs, model_rank = predict_structure(
  File "/home/ubuntu/.local/lib/python3.8/site-packages/colabfold/batch.py", line 441, in predict_structure
  patch_openmm()
  File "/home/ubuntu/.local/lib/python3.8/site-packages/colabfold/batch.py", line 69, in patch_openmm
  from simtk.openmm import app
  ModuleNotFoundError: No module named 'simtk'

Your Environment

cudatoolkit-11.3.1
cudnn-8.2.1
gcc (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.60.02 Driver Version: 510.60.02 CUDA Version: 11.6 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 NVIDIA RTX A6000 On | 00000000:05:00.0 Off | Off |
| 30% 38C P8 18W / 300W | 1MiB / 49140MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+

Additionally, I at first thought my openmm install was incorrect, but running python -m simtk.testInstallation reports all differences within tolerance.

Start with colabfold/batch.py at patch_openmm() and reproduce the reported colabfold_batch command in the listed conda environment. Check the simtk import failure and verify that --amber --use-gpu-relax completes without ModuleNotFoundError.