tf.function property in pytorch

Question

I'm a beginner in pytorch, and I have some functions that are needed to implement in network.

My question is: is there any way like tf.function, or should I use "class(nn.Module)" with variable?

For example, let X be a 10x2 matrix . In pseudo-code:

a = Variable(1.0) 
b = Variable(1.0)
Y = a*X[:,0]**2 + b*X[:,1]

Answer 1

In PyTorch you don't need things like tf.function, you just use normal Python code (because of the dynamic graph).

Please give more detailed example (with code) of what you're trying to do if the above doesn't answer your question.