Skip to content

Conversation

@vpratz
Copy link
Collaborator

@vpratz vpratz commented Dec 9, 2024

Apparently, int32 variables are not transferred to the GPU, leading to problems with XLA. Changing the type declarations to int seems to fix the problem. It occurred at two places:

  • Permutations for the coupling flows.
  • Tracking the training steps for discrete-time consistency models.

Fixes #272.

Apparently, int32 variables are not transferred to the GPU,
leading to problems with XLA. Changing the type declarations
to int seems to fix the problem
@vpratz vpratz requested a review from LarsKue December 9, 2024 15:35
@LarsKue LarsKue requested review from stefanradev93 and removed request for LarsKue December 9, 2024 16:10
@LarsKue
Copy link
Contributor

LarsKue commented Dec 9, 2024

Looks good to me, but I think you would have more insight on this @stefanradev93

@stefanradev93
Copy link
Contributor

Thanks, Valentin! That's an important fix and a good-to-know for the future. :)

@stefanradev93 stefanradev93 merged commit 3f82ce1 into dev Dec 9, 2024
13 checks passed
@stefanradev93 stefanradev93 deleted the fix-gpu-integers branch December 9, 2024 16:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants