-
Notifications
You must be signed in to change notification settings - Fork 68
Description
I've tried af2_interface_metrics.py using /tutorial//halluc_SH3_binder/output/20220104_sh3_r2/trf_relax/00
But, it gave that:
core.chemical.GlobalResidueTypeSet: Total time to initialize 1.19 seconds.
predict_structure(tag_buffer, feature_dict_dict, binderlen_dict, initial_guess_dict, sfd_out, scorefilename)
File "/home1/LJY/rfdesign/RFDesign/scripts/af2_interface_metrics.py", line 433, in predict_structure
prediction_result = jax.vmap(model_runner.apply, in_axes=(None,None,0,0))(model_runner.params,
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/api.py", line 1770, in vmap_f
out_flat = batching.batch(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/linear_util.py", line 165, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 163, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/pjit.py", line 237, in cache_miss
outs, out_flat, out_tree, args_flat = _python_pjit_helper(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/pjit.py", line 180, in _python_pjit_helper
args_flat, _, params, in_tree, out_tree, _ = infer_params_fn(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/api.py", line 443, in infer_params
return pjit.common_infer_params(pjit_info_args, *args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/pjit.py", line 520, in common_infer_params
jaxpr, consts, canonicalized_out_shardings_flat = _pjit_jaxpr(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/linear_util.py", line 301, in memoized_fun
ans = call(fun, *args)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/pjit.py", line 932, in _pjit_jaxpr
jaxpr, global_out_avals, consts = pe.trace_to_jaxpr_dynamic(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/src/profiler.py", line 314, in wrapper
return func(*args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1985, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 2002, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers)
File "/home/users/LJY/.local/lib/python3.8/site-packages/jax/_src/linear_util.py", line 165, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/users/LJY/.local/lib/python3.8/site-packages/haiku/_src/transform.py", line 128, in apply_fn
out, state = f.apply(params, {}, *args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/haiku/_src/transform.py", line 357, in apply_fn
out = f(*args, **kwargs)
jax._src.traceback_util.UnfilteredStackTrace: TypeError: _forward_fn() takes 1 positional argument but 2 were given
The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home1/LJY/rfdesign/RFDesign/scripts/af2_interface_metrics.py", line 599, in
predict_structure(tag_buffer, feature_dict_dict, binderlen_dict, initial_guess_dict, sfd_out, scorefilename)
File "/home1/LJY/rfdesign/RFDesign/scripts/af2_interface_metrics.py", line 433, in predict_structure
prediction_result = jax.vmap(model_runner.apply, in_axes=(None,None,0,0))(model_runner.params,
File "/home/users/LJY/.local/lib/python3.8/site-packages/haiku/_src/transform.py", line 128, in apply_fn
out, state = f.apply(params, {}, *args, **kwargs)
File "/home/users/LJY/.local/lib/python3.8/site-packages/haiku/_src/transform.py", line 357, in apply_fn
out = f(*args, **kwargs)
TypeError: _forward_fn() takes 1 positional argument but 2 were given
What can I do for this error?