The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/root/autodl-tmp/fo/fakeout/avhubert/xunlian.py", line 429, in
app.run(main)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/absl/app.py", line 312, in run
_run_main(main, args)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
sys.exit(main(argv))
File "/root/autodl-tmp/fo/fakeout/avhubert/xunlian.py", line 397, in main
logits = forward_apply(
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/transform.py", line 457, in apply_fn
out = f(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/xunlian.py", line 256, in forward_fn
model_result = module(images=images,
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/mm_embeddings.py", line 240, in call
visual_representation = video_cnn(images, is_training=is_training)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/mm_embeddings.py", line 443, in call
features = self._cnn(images, is_training=is_training)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/tsm_resnet.py", line 329, in call
net = TSMResNetUnit(
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/tsm_resnet.py", line 204, in call
net = TSMResNetBlock(
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/tsm_resnet.py", line 92, in call
preact = self._normalize_fn(preact, is_training=is_training)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/normalization.py", line 143, in
return lambda *a, **k: normalizer_class(**normalizer_kwargs)(*a, **k) # pylint: disable=unnecessary-lambda
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/autodl-tmp/fo/fakeout/avhubert/models/normalization.py", line 85, in call
return super().call(x, is_training,
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 465, in wrapped
out = f(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/contextlib.py", line 75, in inner
return func(*args, **kwds)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/module.py", line 306, in run_interceptors
return bound_method(*args, **kwargs)
File "/root/miniconda3/envs/df/lib/python3.8/site-packages/haiku/_src/batch_norm.py", line 173, in call
mean = jax.lax.pmean(
NameError: unbound axis name: i. The following axis names (e.g. defined by pmap) are available to collective operations: []
Hello，How to solve this problem?
gilikn
