python sample.py --pdb_path ./examples/6zht.pdb --ckpt_path ./weights/mpnn_dataset/epoch=66-step=167499.pt --num_predictions 8 --output_root_dir ./scratch/
[2025-12-13 20:45:53,590] [INFO] [real_accelerator.py:219:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Loaded model weights from ./weights/mpnn_dataset/epoch=66-step=167499.pt
Predictions: 0%| | 0/8 [00:04<?, ?it/s]
Traceback (most recent call last):
File "/mnt/structure/TriFlow-main/sample.py", line 703, in
main()
File "/mnt/structure/TriFlow-main/sample.py", line 595, in main
predicted_sequence = predictor.predict(
^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/sample.py", line 265, in predict
prot_traj, conf = self.interpolant.aa_sample(data, self.model, rigid_frames, aa_init=seq_prior, temp=temp, omit_AA=omit_AA_idx, tied_weights=tied_weights, sample_priority=sample_priority, run_cfg=run_cfg, sample_purity=sample_purity)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/interp/aa_interpolant.py", line 306, in aa_sample
z = model.run_embedder(batch)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/discrete_cond_fm_model.py", line 386, in run_embedder
z = self.triangle_embedder(z, pair_mask.to(dtype=z.dtype), inplace_safe=inplace_safe)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/embedders.py", line 145, in forward
z = tri_att_start_fn(z, pair_mask)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/embedders.py", line 109, in tri_att_start_fn
self.tri_att_start(
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/triangular_attention.py", line 138, in forward
x = self.mha(
^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/primitives.py", line 444, in forward
o = _tri_attention(q, k, v, biases)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/primitives.py", line 476, in _tri_attention
from triflow.utils.tri_kernel.op import TriAttentionFunction
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/op.py", line 18, in
from triflow.utils.tri_kernel.forward import _attention_fwd
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/forward.py", line 22, in
@autotune(
^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/autotune.py", line 241, in decorator
return Autotuner(
^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/autotune.py", line 67, in init
super().init(
TypeError: Autotuner.init() takes from 7 to 10 positional arguments but 13 were given
python sample.py --pdb_path ./examples/6zht.pdb --ckpt_path ./weights/mpnn_dataset/epoch=66-step=167499.pt --num_predictions 8 --output_root_dir ./scratch/
[2025-12-13 20:45:53,590] [INFO] [real_accelerator.py:219:get_accelerator] Setting ds_accelerator to cuda (auto detect)
Loaded model weights from ./weights/mpnn_dataset/epoch=66-step=167499.pt
Predictions: 0%| | 0/8 [00:04<?, ?it/s]
Traceback (most recent call last):
File "/mnt/structure/TriFlow-main/sample.py", line 703, in
main()
File "/mnt/structure/TriFlow-main/sample.py", line 595, in main
predicted_sequence = predictor.predict(
^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/sample.py", line 265, in predict
prot_traj, conf = self.interpolant.aa_sample(data, self.model, rigid_frames, aa_init=seq_prior, temp=temp, omit_AA=omit_AA_idx, tied_weights=tied_weights, sample_priority=sample_priority, run_cfg=run_cfg, sample_purity=sample_purity)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/interp/aa_interpolant.py", line 306, in aa_sample
z = model.run_embedder(batch)
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/discrete_cond_fm_model.py", line 386, in run_embedder
z = self.triangle_embedder(z, pair_mask.to(dtype=z.dtype), inplace_safe=inplace_safe)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/embedders.py", line 145, in forward
z = tri_att_start_fn(z, pair_mask)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/embedders.py", line 109, in tri_att_start_fn
self.tri_att_start(
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/triangular_attention.py", line 138, in forward
x = self.mha(
^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/miniconda/envs/torch11/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
return forward_call(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/primitives.py", line 444, in forward
o = _tri_attention(q, k, v, biases)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/multimer_model/primitives.py", line 476, in _tri_attention
from triflow.utils.tri_kernel.op import TriAttentionFunction
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/op.py", line 18, in
from triflow.utils.tri_kernel.forward import _attention_fwd
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/forward.py", line 22, in
@autotune(
^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/autotune.py", line 241, in decorator
return Autotuner(
^^^^^^^^^^
File "/mnt/structure/TriFlow-main/triflow/utils/tri_kernel/autotune.py", line 67, in init
super().init(
TypeError: Autotuner.init() takes from 7 to 10 positional arguments but 13 were given