свойство tf.function в pytorch

Question

Я новичок в pytorch, и у меня есть некоторые функции, которые необходимо реализовать в сети.

Мой вопрос: есть ли способ, подобный tf.function , или я должен использовать «class (nn.Module)» с переменной?

Например, пусть X будет матрицей 10x2. В псевдокоде:

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

Answer 1

В PyTorch вам не нужны такие вещи, как tf.function, вы просто используете обычный код Python (из-за динамического графика).

Пожалуйста, приведите более подробный пример (с кодом) того, что выпытаясь сделать, если вышеупомянутое не отвечает на ваш вопрос.