Replies: 1 comment
-
Hey @cisprague, can you verify that |
Beta Was this translation helpful? Give feedback.
0 replies
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hi! I am a big fan of NNX so far, great job!
I am trying to make a "training" module that includes everything (model, dataset, metrics, etc.), but I find that
nnx.jit
ing thetrain_step
method when usingself
is much slower. As a result, I resorted to using staticmethod. Are there some best practices on how to do this properly? Please see the code below (not the full code, but I hope it gets the point across).Thank you!
Beta Was this translation helpful? Give feedback.
All reactions