Save result from jax.local_device_count

google / nerfies

This is the code for Deformable Neural Radiance Fields, a.k.a. Nerfies.

https://nerfies.github.io

Apache License 2.0

1.59k stars 215 forks source link

Save result from jax.local_device_count #47

Closed JamesPerlman closed 2 years ago

JamesPerlman commented 2 years ago

jax.local_device_count() result is discarded here, but I do not believe there is a way to reach this codepath (currently there are no usages of utils.shard(...) missing the device_count arg)

keunhong commented 2 years ago

Strange, not sure how that happened. Thanks for the fix!