When i run the command:
python -m jaxnerf.train --data_dir=./nerf_data --train_dir=./checkpoints --config=configs/objectron_small
shown in the Objectron_NeRF_Tutorial.ipynb i get the following output:
I0419 18:05:48.879847 20420 xla_bridge.py:161] Remote TPU is not linked into jax; skipping remote TPU.
I0419 18:05:48.879847 20420 xla_bridge.py:328] Unable to initialize backend 'tpu_driver': Could not initialize backend 'tpu_driver'
I0419 18:05:48.971926 20420 xla_bridge.py:328] Unable to initialize backend 'rocm': NOT_FOUND: Could not find registered platform with name: "rocm". Available platform names are: Interpreter CUDA Host
I0419 18:05:48.975028 20420 xla_bridge.py:328] Unable to initialize backend 'tpu': module 'jaxlib.xla_extension' has no attribute 'get_tpu_client'
2023-04-19 18:05:49.117062: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc:57] cuLinkAddData fails. This is usually caused by stale driver version.
2023-04-19 18:05:49.117268: E external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc:1334] The CUDA linking API did not work. Please use XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1 to bypass it, but expect to get longer compilation time due to the lack of multi-threading.
Traceback (most recent call last):
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\runpy.py", line 194, in _run_module_as_main
return _run_code(code, main_globals, None,
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\runpy.py", line 87, in _run_code
exec(code, run_globals)
File "C:\Users\user\Downloads\TFMRuben\jaxnerf\train.py", line 269, in
app.run(main)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\absl\app.py", line 308, in run
_run_main(main, args)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\absl\app.py", line 254, in _run_main
sys.exit(main(argv))
File "C:\Users\user\Downloads\TFMRuben\jaxnerf\train.py", line 116, in main
rng = random.PRNGKey(20200823)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\random.py", line 125, in PRNGKey
key = prng.seed_with_impl(impl, seed)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\prng.py", line 233, in seed_with_impl
return PRNGKeyArray(impl, impl.seed(seed))
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\prng.py", line 272, in threefry_seed
lax.shift_right_logical(seed_arr, lax_internal._const(seed_arr, 32)))
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\lax\lax.py", line 487, in shift_right_logical
return shift_right_logical_p.bind(x, y)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax\core.py", line 324, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax\core.py", line 327, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax\core.py", line 684, in process_primitive
return primitive.impl(*tracers, **params)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 99, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\util.py", line 220, in wrapper
return cached(config._trace_context(), *args, **kwargs)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\util.py", line 213, in cached
return f(*args, **kwargs)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 164, in xla_primitive_callable
compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 248, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 827, in compile
self._executable = XlaCompiledComputation.from_xla_computation(
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 934, in from_xla_computation
compiled = compile_or_get_cached(backend, xla_computation, options,
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 899, in compile_or_get_cached
return backend_compile(backend, computation, compile_options, host_callbacks)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "C:\Users\user\anaconda3\envs\jaxnerf2\lib\site-packages\jax_src\dispatch.py", line 843, in backend_compile
return backend.compile(built_c, compile_options=options)
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: no kernel image is available for execution on the device
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_asm_compiler.cc(60): 'status'
google-research-datasets
