JIT with lambda functions

[tan-liam]

When initializing an nn module, I pass in a lambda function and save it as self.activation. It is possible to use this activation and use partial jit, barring the activation only from being Jitted?

JIT is usually used in the outer layer.

Is it possible to use JIT when i use a lambda function, passed in the lambda function?

e.g.

x = self.activation(x), where self.activation is a lambda function.

[cgarciae]

Hey @tan-liam, if you are inside jit, all other inner jits have no effect so I think you will have no problem.