Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[Bug] Example is broken on GPU #238

Open
CForgie opened this issue Jun 24, 2022 · 1 comment
Open

[Bug] Example is broken on GPU #238

CForgie opened this issue Jun 24, 2022 · 1 comment
Labels
bug Something isn't working

Comments

@CForgie
Copy link

CForgie commented Jun 24, 2022

The example listed in https://poets-ai.github.io/elegy/getting-started/high-level-api/ doesn't work on GPU with the latest version of elegy - when calling model.fit training starts to take place but stops short of completing one epoch and hangs. Downgrading to 0.8.4 seems to fix this.

Library Info

elegy 0.8.6
flax 0.4.2
jax 0.3.13
jax-metrics 0.1.2
jaxlib 0.3.10
optax 0.1.2
treex 0.6.10

CUDA Version: 11.6
Tesla T4

Running ami: Deep Learning Base AMI (Amazon Linux 2) Version 53*
@CForgie CForgie added the bug Something isn't working label Jun 24, 2022
@CForgie
Copy link
Author

CForgie commented Jun 24, 2022

Seems similar to #234

@CForgie CForgie changed the title Example is broken on GPU [Bug] [Bug] Example is broken on GPU Jun 24, 2022
@jiyuuchc jiyuuchc mentioned this issue Aug 21, 2022
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@CForgie